left only answers

pipeline.py

@@ -14,13 +14,16 @@ class PreTrainedPipeline():
         self.tokenizer = transformers.AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
 
     def __call__(self, inputs: str) -> List[Dict]:
-
+        # Get input tokens
         text = inputs + self.tokenizer.eos_token
         start_tokens = self.tokenizer.convert_ids_to_tokens(self.tokenizer.encode(text))
-
+        # generate
         results = self.generator.generate_batch([start_tokens])
         output = results[0].sequences[0]
-
-        generated_text = self.tokenizer.decode(self.tokenizer.convert_tokens_to_ids(output))
+        # 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}]