Output cleanup

If the answer has context and reponse tags; only return the repsonse or a message that you can't help.

handler.py

@@ -93,10 +93,21 @@ class EndpointHandler:
             output_tokens[0]
         )
 
+        def substring_after(s, delim):
+            return s.partition(delim)[2]
+
         answer = generated_text.split('>>END<<')[0].split('>>ANSWER<<:')[1].strip()
 
+        if "CONTEXT:" in answer:
+            if "RESPONSE:" in answer:
+                answerclean = answer.partition("RESPONSE:")[2]
+            else:
+                answerclean = "I'm sorry. I'm not able to help with this tender topic."
+        else:
+            answerclean = answer
+
         #prediction = {'generated_text': answer.replace('<|endoftext|>', ''), 'time': f"{(end-start):.2f} s"} 
-        prediction = {'generated_text': answer.replace('<|endoftext|>', ''), 'time': f"{(end-start):.2f} s", 'complete_ouput': generated_text}
+        prediction = {'generated_text': answerclean.replace('<|endoftext|>', ''), 'time': f"{(end-start):.2f} s", 'complete_ouput': generated_text}
 
         result = []
         result.append(prediction)