Use opt for generation

pipeline.py

@@ -9,9 +9,10 @@ import os
 class PreTrainedPipeline():
     def __init__(self, path: str):
         # Init DialoGPT
-        dialogpt_path = os.path.join(path, "dialogpt")
-        self.generator = ctranslate2.Generator(dialogpt_path, device="cpu", compute_type="int8")
-        self.tokenizer = transformers.AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+        self.eos_token = "\n"
+        dialogpt_path = os.path.join(path, "opt")
+        self.generator = ctranslate2.Generator(dialogpt_path, device="cpu", compute_type="float")
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained("facebook/opt-350m")
         # Init M2M100
         m2m100_path = os.path.join(path, "m2m100")
         self.translator = ctranslate2.Translator(m2m100_path, device="cpu", compute_type="int8")
@@ -29,17 +30,22 @@ class PreTrainedPipeline():
 
     def dialogpt(self, inputs: str) -> str:
         # Get input tokens
-        text = inputs + self.tokenizer.eos_token
+        text = inputs + self.eos_token
         start_tokens = self.tokenizer.convert_ids_to_tokens(self.tokenizer.encode(text))
         # generate
-        results = self.generator.generate_batch([start_tokens])
+        results = self.generator.generate_batch([start_tokens], max_length=50, repetition_penalty=1.2)
         output = results[0].sequences[0]
         # left only answers
         tokens = self.tokenizer.convert_tokens_to_ids(output)
-        eos_index = tokens.index(self.tokenizer.eos_token_id)
-        answer_tokens = tokens[eos_index+1:]
-        generated_text = self.tokenizer.decode(answer_tokens)
-        return generated_text
+        generated_text = self.tokenizer.decode(tokens)
+        eos_index = self.index_last(generated_text, self.eos_token)
+        answer_text = generated_text[eos_index+1:]
+        return answer_text
+
+    @staticmethod
+    def index_last(li: str, char: str):
+        idx = len(li) - 1 - li[::-1].index(char)
+        return idx 
 
     def m2m100(self, inputs: str, from_lang: str, to_lang: str) -> str:
         self.m2m100_tokenizer.src_lang = from_lang