Upload handler.py

handler.py

@@ -22,7 +22,6 @@ class EndpointHandler():
         self.model = PeftModel.from_pretrained(model, path)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        LOGGER.info(f"Received data: {data}")
         # Get inputs
         # Extract required parameters from data
         message = data.get("message")
@@ -127,18 +126,25 @@ def generate(
 
     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)
         for end_sequence in end_sequences:
             if end_sequence in generated_text:
-                streaming = False
                 generated_text = generated_text.replace(end_sequence, "")
+                return parse(generated_text, conclusions, end_sequences)
+
+def parse(generated_text: str, conclusions: list[tuple[str, str]], end_sequences: list[str]) -> dict:
+    # Initialize variables
+    conclusion_found = None
+    context_numbers = []
+
+    # Remove end sequences and clean the text
+    for end_sequence in end_sequences:
+        generated_text = generated_text.replace(end_sequence, "")
+    generated_text = generated_text.strip()
 
     # Check for conclusion keys in the generated text
     if conclusions:
@@ -153,7 +159,7 @@ def generate(
     context_numbers = [int(match.strip("[]")) for match in context_matches]
 
     return {
-        "generated_text": generated_text.strip(),
+        "generated_text": generated_text,
         "conclusion": conclusion_found,
         "context": context_numbers
-    }
+    }