INNER CODE UNIT · Python

SimCTGDialogue

yxuansu/SimCTG · dialogue_generation/simctgdialogue.py:20

class SimCTGDialogue(nn.Module):
    def __init__(self, model_name, eos_token, pad_token):
        super(SimCTGDialogue, self).__init__()
        from transformers import AutoTokenizer, GPT2LMHeadModel
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.vocab_size = len(self.tokenizer)
        self.model = GPT2LMHeadModel.from_pretrained(model_name)
        self.embed_dim = self.model.config.hidden_size
        if pad_token in self.tokenizer.vocab:
            print ('PAD token exists.')
        else:
            print ('Add PAD token to the tokenizer.')
            print ('Original vocabulary size is {}'.format(len(self.tokenizer)))
            self.tokenizer.add_tokens([pad_token])
            print ('Vocabulary size after extension is {}'.format(len(self.tokenizer)))
            assert len(self.tokenizer.convert_tokens_to_ids([pad_token])) == 1
            self.model.resize_token_embeddings(len(self.tokenizer)) 
        self.pad_token_id = self.tokenizer.convert_tokens_to_ids([pad_token])[0]

View source record →