from transformers import T5TokenizerFast, T5ForConditionalGeneration, GenerationConfig from model import Model class T5(Model): def __init__(self, name:str='T5', model_dir:str='./models/pko_t5_COMU_patience10', max_input_length:int=64, max_target_length:int=64 ): self.name = name self.model = T5ForConditionalGeneration.from_pretrained(model_dir) self.tokenizer = T5TokenizerFast.from_pretrained(model_dir) self.gen_config = GenerationConfig.from_pretrained(model_dir, 'gen_config.json') self.max_input_length = max_input_length self.max_target_length = max_target_length self.INPUT_FORMAT = 'qa question: ' # add tokens self.tokenizer.add_tokens(["#화자#", "#청자#", "#(남자)청자#", "#(남자)화자#", "#(여자)청자#", "(여자)화자"]) self.model.resize_token_embeddings(len(self.tokenizer)) self.model.config.max_length = max_target_length self.tokenizer.model_max_length = max_target_length def generate(self, inputs): inputs = self.INPUT_FORMAT.replace("", inputs) input_ids = self.tokenizer(inputs, max_length=self.max_input_length, truncation=True, return_tensors="pt") output_tensor = self.model.generate(**input_ids, generation_config=self.gen_config) output_ids = self.tokenizer.batch_decode(output_tensor, skip_special_tokens=True, clean_up_tokenization_spaces=True) outputs = str(output_ids) outputs = outputs.replace('[', '').replace(']', '').replace("'", '').replace("'", '') return outputs