from transformers import T5TokenizerFast, T5ForConditionalGeneration

class T5:
    def __init__(self,
            model_dir:str='./model/pko_t5_COMU_patience10',
            max_input_length:int=64,
            max_target_length:int=64,
            prefix:str='qa question: '
            ):
        self.model = T5ForConditionalGeneration.from_pretrained(model_dir)
        self.tokenizer = T5TokenizerFast.from_pretrained(model_dir)
        self.max_input_length = max_input_length
        self.max_target_length = max_target_length
        self.prefix = prefix
        
        # 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 chat(self, inputs):
        inputs = [self.prefix + inputs]
        input_ids = self.tokenizer(inputs, max_length=self.max_input_length, truncation=True, return_tensors="pt")
        output_tensor = self.model.generate(**input_ids, num_beams=2, do_sample=True, min_length=10, max_length=self.max_target_length, no_repeat_ngram_size=2) #repetition_penalty=2.5
        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