model/dialogue/hmnet_model.py

from model.base_model import SummModel
import argparse
import os
import torch
import gzip
import json
from model.third_party.HMNet.Models.Trainers.HMNetTrainer import HMNetTrainer
from model.third_party.HMNet.Utils.Arguments import Arguments

import spacy

nlp = spacy.load("en_core_web_sm", disable=["parser"])
# tagger = nlp.get_pipe('tagger')
# ner = nlp.get_pipe('ner')
# POS = {w: i for i, w in enumerate([''] + list(tagger.labels))}
# ENT = {w: i for i, w in enumerate([''] + list(ner.move_names))}
# These two dicts are adapted from SpaCy 2.3.1, since HMNet's embedding for POS and ENT is fixed
POS = {
    "": 0,
    "$": 1,
    "''": 2,
    ",": 3,
    "-LRB-": 4,
    "-RRB-": 5,
    ".": 6,
    ":": 7,
    "ADD": 8,
    "AFX": 9,
    "CC": 10,
    "CD": 11,
    "DT": 12,
    "EX": 13,
    "FW": 14,
    "HYPH": 15,
    "IN": 16,
    "JJ": 17,
    "JJR": 18,
    "JJS": 19,
    "LS": 20,
    "MD": 21,
    "NFP": 22,
    "NN": 23,
    "NNP": 24,
    "NNPS": 25,
    "NNS": 26,
    "PDT": 27,
    "POS": 28,
    "PRP": 29,
    "PRP$": 30,
    "RB": 31,
    "RBR": 32,
    "RBS": 33,
    "RP": 34,
    "SYM": 35,
    "TO": 36,
    "UH": 37,
    "VB": 38,
    "VBD": 39,
    "VBG": 40,
    "VBN": 41,
    "VBP": 42,
    "VBZ": 43,
    "WDT": 44,
    "WP": 45,
    "WP$": 46,
    "WRB": 47,
    "XX": 48,
    "_SP": 49,
    "``": 50,
}
ENT = {
    "": 0,
    "B-ORG": 1,
    "B-DATE": 2,
    "B-PERSON": 3,
    "B-GPE": 4,
    "B-MONEY": 5,
    "B-CARDINAL": 6,
    "B-NORP": 7,
    "B-PERCENT": 8,
    "B-WORK_OF_ART": 9,
    "B-LOC": 10,
    "B-TIME": 11,
    "B-QUANTITY": 12,
    "B-FAC": 13,
    "B-EVENT": 14,
    "B-ORDINAL": 15,
    "B-PRODUCT": 16,
    "B-LAW": 17,
    "B-LANGUAGE": 18,
    "I-ORG": 19,
    "I-DATE": 20,
    "I-PERSON": 21,
    "I-GPE": 22,
    "I-MONEY": 23,
    "I-CARDINAL": 24,
    "I-NORP": 25,
    "I-PERCENT": 26,
    "I-WORK_OF_ART": 27,
    "I-LOC": 28,
    "I-TIME": 29,
    "I-QUANTITY": 30,
    "I-FAC": 31,
    "I-EVENT": 32,
    "I-ORDINAL": 33,
    "I-PRODUCT": 34,
    "I-LAW": 35,
    "I-LANGUAGE": 36,
    "L-ORG": 37,
    "L-DATE": 38,
    "L-PERSON": 39,
    "L-GPE": 40,
    "L-MONEY": 41,
    "L-CARDINAL": 42,
    "L-NORP": 43,
    "L-PERCENT": 44,
    "L-WORK_OF_ART": 45,
    "L-LOC": 46,
    "L-TIME": 47,
    "L-QUANTITY": 48,
    "L-FAC": 49,
    "L-EVENT": 50,
    "L-ORDINAL": 51,
    "L-PRODUCT": 52,
    "L-LAW": 53,
    "L-LANGUAGE": 54,
    "U-ORG": 55,
    "U-DATE": 56,
    "U-PERSON": 57,
    "U-GPE": 58,
    "U-MONEY": 59,
    "U-CARDINAL": 60,
    "U-NORP": 61,
    "U-PERCENT": 62,
    "U-WORK_OF_ART": 63,
    "U-LOC": 64,
    "U-TIME": 65,
    "U-QUANTITY": 66,
    "U-FAC": 67,
    "U-EVENT": 68,
    "U-ORDINAL": 69,
    "U-PRODUCT": 70,
    "U-LAW": 71,
    "U-LANGUAGE": 72,
    "O": 73,
}


class HMNetModel(SummModel):
    # static variables
    model_name = "HMNET"
    is_extractive = False
    is_neural = True
    is_dialogue_based = True

    def __init__(
        self,
        min_gen_length: int = 10,
        max_gen_length: int = 300,
        beam_width: int = 6,
        **kwargs,
    ):
        """
        Create a summarization model with HMNet backbone. In the default setting, the inference speed will be
        10s/sample (on one GPU), however, if one can tune these three parameters properly, e.g. min_gen_length=10,
        max_gen_length=100, and beam_width=2, the inference speed will increase to 2s/sample (on one GPU).

        Args:
          min_gen_length (int): minimum generation length of the decoder
          max_gen_length (int): maximum generation length of the decoder
          beam_width (int): width of the beam when doing beam search in the decoding process
          kwargs: the other valid parameters. The valid parameters can be found in
              model/dialogue/hmnet/config/dialogue.conf . You can use either lower case or upper case for parameter
              name. The valid parameter name is one of the following args, however, we do not encourage you to modify
               them, since some unexpected, untested errors might be triggered:
              ['MODEL', 'TASK', 'CRITERION', 'SEED', 'MAX_NUM_EPOCHS', 'EVAL_PER_UPDATE_NUM'
              , 'UPDATES_PER_EPOCH', 'OPTIMIZER', 'START_LEARNING_RATE', 'LR_SCHEDULER', 'WARMUP_STEPS',
              'WARMUP_INIT_LR', 'WARMUP_END_LR', 'GRADIENT_ACCUMULATE_STEP', 'GRAD_CLIPPING', 'USE_REL_DATA_PATH',
              'TRAIN_FILE', 'DEV_FILE', 'TEST_FILE', 'ROLE_DICT_FILE', 'MINI_BATCH', 'MAX_PADDING_RATIO',
              'BATCH_READ_AHEAD', 'DOC_SHUFFLE_BUF_SIZE', 'SAMPLE_SHUFFLE_BUFFER_SIZE', 'BATCH_SHUFFLE_BUFFER_SIZE',
              'MAX_TRANSCRIPT_WORD', 'MAX_SENT_LEN', 'MAX_SENT_NUM', 'DROPOUT', 'VOCAB_DIM', 'ROLE_SIZE', 'ROLE_DIM',
              'POS_DIM', 'ENT_DIM', 'USE_ROLE', 'USE_POSENT', 'USE_BOS_TOKEN', 'USE_EOS_TOKEN',
              'TRANSFORMER_EMBED_DROPOUT', 'TRANSFORMER_RESIDUAL_DROPOUT', 'TRANSFORMER_ATTENTION_DROPOUT',
              'TRANSFORMER_LAYER', 'TRANSFORMER_HEAD', 'TRANSFORMER_POS_DISCOUNT', 'PRE_TOKENIZER',
              'PRE_TOKENIZER_PATH', 'PYLEARN_MODEL', 'EXTRA_IDS', 'BEAM_WIDTH', 'EVAL_TOKENIZED', 'EVAL_LOWERCASE',
              'MAX_GEN_LENGTH', 'MIN_GEN_LENGTH', 'NO_REPEAT_NGRAM_SIZE']

        Return an instance of HMNet model for dialogue summarization.
        """
        super(HMNetModel, self).__init__()
        self.root_path = self._get_root()

        # we leave the most influential params with prompt and the others as hidden kwargs
        kwargs["MIN_GEN_LENGTH"] = min_gen_length
        kwargs["MAX_GEN_LENGTH"] = max_gen_length
        kwargs["BEAM_WIDTH"] = beam_width
        self.opt = self._parse_args(kwargs)
        self.model = HMNetTrainer(self.opt)

    def _get_root(self):
        root_path = os.getcwd()
        while "model" not in os.listdir(root_path):
            root_path = os.path.dirname(root_path)
        root_path = os.path.join(root_path, "model/dialogue")
        return root_path

    def _parse_args(self, kwargs):
        parser = argparse.ArgumentParser(
            description="HMNet: Pretrain or fine-tune models for HMNet model."
        )
        parser.add_argument(
            "--command", default="evaluate", help="Command: train/evaluate"
        )
        parser.add_argument(
            "--conf_file",
            default=os.path.join(self.root_path, "hmnet/config/dialogue.conf"),
            help="Path to the BigLearn conf file.",
        )
        parser.add_argument(
            "--PYLEARN_MODEL", help="Overrides this option from the conf file."
        )
        parser.add_argument(
            "--master_port", help="Overrides this option default", default=None
        )
        parser.add_argument("--cluster", help="local, philly or aml", default="local")
        parser.add_argument(
            "--dist_init_path", help="Distributed init path for AML", default="./tmp"
        )
        parser.add_argument(
            "--fp16",
            action="store_true",
            help="Whether to use 16-bit float precision instead of 32-bit",
        )
        parser.add_argument(
            "--fp16_opt_level",
            type=str,
            default="O1",
            help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
            "See details at https://nvidia.github.io/apex/amp.html",
        )
        parser.add_argument("--no_cuda", action="store_true", help="Disable cuda.")
        parser.add_argument(
            "--config_overrides",
            help="Override parameters on config, VAR=val;VAR=val;...",
        )

        cmdline_args = parser.parse_args()
        command = cmdline_args.command
        conf_file = cmdline_args.conf_file
        conf_args = Arguments(conf_file)
        opt = conf_args.readArguments()

        if cmdline_args.config_overrides:
            for config_override in cmdline_args.config_overrides.split(";"):
                config_override = config_override.strip()
                if config_override:
                    var_val = config_override.split("=")
                    assert (
                        len(var_val) == 2
                    ), f"Config override '{var_val}' does not have the form 'VAR=val'"
                    conf_args.add_opt(opt, var_val[0], var_val[1], force_override=True)

        opt["cuda"] = torch.cuda.is_available() and not cmdline_args.no_cuda
        opt["confFile"] = conf_file
        if "datadir" not in opt:
            opt["datadir"] = os.path.dirname(
                conf_file
            )  # conf_file specifies where the data folder is
        opt["basename"] = os.path.basename(
            conf_file
        )  # conf_file specifies where the name of save folder is
        opt["command"] = command

        # combine cmdline_args into opt dictionary
        for key, val in cmdline_args.__dict__.items():
            # if val is not None and key not in ['command', 'conf_file']:
            if val is not None:
                opt[key] = val

        # combine kwargs into opt dictionary (we allow lower case)
        for key, val in kwargs.items():
            valid_keys = [x for x in opt.keys() if x.upper() == x]
            if key.upper() not in valid_keys:
                print("WARNING: {} is not a valid key in HMNet.".format(key))
                print("The valid keys are:", valid_keys)
                continue
            if val is not None:
                opt[key.upper()] = val

        return opt

    def summarize(self, corpus, queries=None):
        print(f"HMNet model: processing document of {corpus.__len__()} samples")
        # transform the original dataset to "dialogue" input
        # we only use test set path for evaluation
        data_folder = os.path.join(
            os.path.dirname(self.opt["datadir"]),
            "ExampleRawData/meeting_summarization/AMI_proprec/test",
        )

        self._create_datafolder(data_folder)
        self._preprocess(corpus, data_folder)

        # return self.model.eval()
        results = self._evaluate()

        return results

    def _evaluate(self):
        if self.opt["rank"] == 0:
            self.model.log("-----------------------------------------------")
            self.model.log("Evaluating model ... ")

        self.model.set_up_model()

        eval_dataset = "test"
        batch_generator_eval = self.model.get_batch_generator(eval_dataset)
        predictions = self._eval_batches(
            self.model.module, batch_generator_eval, self.model.saveFolder, eval_dataset
        )

        return predictions

    def _eval_batches(self, module, dev_batches, save_folder, label=""):
        max_sent_len = int(self.opt["MAX_GEN_LENGTH"])

        print("Decoding current model ... \nSaving folder is {}".format(save_folder))
        print("Each sample will cost about 10 second.")
        import time

        start_time = time.time()
        predictions = []  # prediction of tokens from model
        if not isinstance(module.tokenizer, list):
            decoder_tokenizer = module.tokenizer
        elif len(module.tokenizer) == 1:
            decoder_tokenizer = module.tokenizer[0]
        elif len(module.tokenizer) == 2:
            decoder_tokenizer = module.tokenizer[1]
        else:
            assert False, "len(module.tokenizer) > 2"

        with torch.no_grad():
            for j, dev_batch in enumerate(dev_batches):
                for b in dev_batch:
                    if torch.is_tensor(dev_batch[b]):
                        dev_batch[b] = dev_batch[b].to(self.opt["device"])

                beam_search_res = module(
                    dev_batch, beam_search=True, max_sent_len=max_sent_len
                )
                pred = [
                    [t[0] for t in x] if len(x) > 0 else [[]] for x in beam_search_res
                ]
                predictions.extend(
                    [
                        [
                            self._convert_tokens_to_string(decoder_tokenizer, tt)
                            for tt in t
                        ]
                        for t in pred
                    ]
                )

                if (
                    "DEBUG" in self.opt and j >= 10
                ) or j >= self.model.task.evaluator.eval_batches_num:
                    # in debug mode (decode first 10 batches) ortherwise decode first self.eval_batches_num bathes
                    break

        top1_predictions = [x[0] for x in predictions]

        print("Total time for inference:", time.time() - start_time)
        return top1_predictions

    def _convert_tokens_to_string(self, tokenizer, tokens):
        if "EVAL_TOKENIZED" in self.opt:
            tokens = [t for t in tokens if t not in tokenizer.all_special_tokens]
        if "EVAL_LOWERCASE" in self.opt:
            tokens = [t.lower() for t in tokens]
        if "EVAL_TOKENIZED" in self.opt:
            return " ".join(tokens)
        else:
            return tokenizer.decode(
                tokenizer.convert_tokens_to_ids(tokens), skip_special_tokens=True
            )

    def _preprocess(self, corpus, test_path):
        samples = []
        for i, sample in enumerate(corpus):
            new_sample = {"id": i, "meeting": [], "summary": []}
            if isinstance(sample, str):
                raise RuntimeError(
                    "Error: the input of HMNet should be dialogues, rather than documents."
                )

            # add all the turns one by one
            for turn in sample:
                turn = [x.strip() for x in turn.split(":")]
                if len(turn) < 2:
                    continue
                tokenized_turn = nlp(turn[1])
                # In case we can't find proper entity in move_names
                ent_id = []
                pos_id = []
                for token in tokenized_turn:
                    ent = (
                        token.ent_iob_ + "-" + token.ent_type_
                        if token.ent_iob_ != "O"
                        else "O"
                    )
                    ent_id.append(ENT[ent] if ent in ENT else ENT[""])

                    pos = token.tag_
                    pos_id.append(POS[pos] if pos in POS else POS[""])

                new_sample["meeting"].append(
                    {
                        "speaker": turn[0],
                        "role": "",
                        "utt": {
                            "word": [str(token) for token in tokenized_turn],
                            "pos_id": pos_id,
                            "ent_id": ent_id,
                        },
                    }
                )
            new_sample["summary"].append(
                "This is a dummy summary. HMNet will filter out the sample w/o summary!"
            )
            samples.append(new_sample)
            # save to the gzip
            file_path = os.path.join(test_path, "split_{}.jsonl.gz".format(i))
            with gzip.open(file_path, "wt", encoding="utf-8") as file:
                file.write(json.dumps(new_sample))

    def _clean_datafolder(self, data_folder):
        for name in os.listdir(data_folder):
            name = os.path.join(data_folder, name)
            if ".gz" in name:
                os.remove(name)

    def _create_datafolder(self, data_folder):
        if os.path.exists(data_folder):
            self._clean_datafolder(data_folder)
        else:
            os.makedirs(data_folder)
        with open(
            os.path.join(os.path.dirname(data_folder), "test_ami.json"),
            "w",
            encoding="utf-8",
        ) as file:
            json.dump(
                [
                    {
                        "source": {
                            "dataset": "../ExampleRawData/meeting_summarization/AMI_proprec/test/"
                        },
                        "task": "meeting",
                        "name": "ami",
                    }
                ],
                file,
            )

        with open(
            os.path.join(
                os.path.dirname(os.path.dirname(data_folder)), "role_dict_ext.json"
            ),
            "w",
        ) as file:
            json.dump({}, file)

    @classmethod
    def show_capability(cls) -> None:
        basic_description = cls.generate_basic_description()
        more_details = (
            "A HMNet model finetuned on CNN-DM dataset for summarization.\n\n"
            "Strengths:\n - High performance on dialogue summarization task.\n\n"
            "Weaknesses:\n - Not suitable for datasets other than dialogues.\n\n"
            "Initialization arguments:\n "
            " - `corpus`: Unlabelled corpus of documents.\n"
        )
        print(f"{basic_description} \n {'#' * 20} \n {more_details}")