Update pipeline.py

pipeline.py

@@ -3,8 +3,9 @@ from transformers import Text2TextGenerationPipeline, AutoModelForSeq2SeqLM, Aut
 
 def get_markers_for_model():
     special_tokens_constants = Namespace() 
-    special_tokens_constants.separator_different_qa = "&&"
-    special_tokens_constants.separator_output_question_answer = "? "
+    special_tokens_constants.separator_different_qa = "&&&"
+    special_tokens_constants.separator_output_question_answer = "SSEEPP"
+    special_tokens_constants.source_prefix = "qa: "
     return special_tokens_constants
 
 def load_trained_model(name_or_path):
@@ -21,10 +22,19 @@ class QADiscourse_Pipeline(Text2TextGenerationPipeline):
     
 
     def preprocess(self, inputs):
-        # Here, inputs is string or list of strings; apply string postprocessing
-        return super().preprocess(inputs)
+        if isinstance(inputs, str):
+            processed_inputs = self._preprocess_string(inputs)
+        elif hasattr(inputs, "__iter__"):
+            processed_inputs = [self._preprocess_string(s) for s in inputs]
+        else:
+            raise ValueError("inputs must be str or Iterable[str]")
+        # Now pass to super.preprocess for tokenization
+        return super().preprocess(processed_inputs)
     
-
+    def _preprocess_string(self, seq: str) -> str:
+        seq = self.special_tokens.source_prefix + seq
+        print(seq)
+        return seq
     
     def _forward(self, *args, **kwargs):
         outputs = super()._forward(*args, **kwargs)
@@ -36,7 +46,9 @@ class QADiscourse_Pipeline(Text2TextGenerationPipeline):
         seperated_qas = self._split_to_list(predictions)
         qas = []
         for qa_pair in seperated_qas:
-            qas.append(self._postrocess_qa(qa_pair))
+            post_process = self._postrocess_qa(qa_pair) # if the prediction isn't a valid QA
+            if post_process is not None:
+                qas.append(post_process)
         return qas
 
     def _split_to_list(self, output_seq: str) -> list:
@@ -48,7 +60,6 @@ class QADiscourse_Pipeline(Text2TextGenerationPipeline):
         if self.special_tokens.separator_output_question_answer in seq:
             question, answer = seq.split(self.special_tokens.separator_output_question_answer)
         else:
-            print("invalid format: no separator between question and answer found...")
             return None
         return {"question": question, "answer": answer}
     
@@ -59,5 +70,4 @@ if __name__ == "__main__":
     res2 = pipe(["I don't like chocolate, but I like cookies.",
                  "I dived in the sea easily"], num_beams=10)
     print(res1)
-    print(res2)
-    
+    print(res2)