Upload handler.py

handler.py

@@ -64,7 +64,7 @@ class EndpointHandler():
         # Postprocess
         prediction = output
         LOGGER.info(f"Generated text: {prediction}")
-        return {"generated_text": prediction}
+        return prediction
 
 def generate(
     tokenizer,
@@ -121,31 +121,39 @@ def generate(
         top_k=top_k,
         temperature=temperature,
         num_beams=1,
-        repetition_penalty=repetition_penalty,
+        repetition_penalty=repetition_penalty
     )
     model.generate(**generate_kwargs)
 
     outputs = []
+    generated_text = ""
+    streaming = True
+    conclusion_found = None
+    context_numbers = []
     for text in streamer:
+        if not streaming:
+            break
         outputs.append(text)
         generated_text = "".join(outputs)
-        conclusion_found = None
-        context_numbers = []
-
-        # Check for conclusion keys in the generated text
-        if conclusions:
-            for conclusion_key, _ in conclusions:
-                if conclusion_key in generated_text:
-                    conclusion_found = conclusion_key
-                    break
-
-        # Extract context numbers from the generated text
-        context_pattern = r"\[\d+\]"
-        context_matches = re.findall(context_pattern, generated_text)
-        context_numbers = [int(match.strip("[]")) for match in context_matches]
-
-        return {
-            "generated_text": generated_text.strip(),
-            "conclusion": conclusion_found,
-            "context": context_numbers
-        }
+        for end_sequence in end_sequences:
+            if end_sequence in generated_text:
+                streaming = False
+                generated_text = generated_text.replace(end_sequence, "")
+
+    # Check for conclusion keys in the generated text
+    if conclusions:
+        for conclusion_key, _ in conclusions:
+            if conclusion_key in generated_text:
+                conclusion_found = conclusion_key
+                break
+
+    # Extract context numbers from the generated text
+    context_pattern = r"\[\d+\]"
+    context_matches = re.findall(context_pattern, generated_text)
+    context_numbers = [int(match.strip("[]")) for match in context_matches]
+
+    return {
+        "generated_text": generated_text.strip(),
+        "conclusion": conclusion_found,
+        "context": context_numbers
+    }