INNER CODE UNIT · Python
__init__
yxuansu/SimCTG · dialogue_generation/simctgdialogue.py:21
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]
self.vocab_size = len(self.tokenizer)