Upload handler.py

handler.py

@@ -79,13 +79,14 @@ def generate(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-) -> str:
+    end_sequences: list[str] = ["[INST]", "[/INST]", "\n"]
+) -> dict:
     
     # Check if the system_prompt is provided, else construct it from instruction, conclusions, and context
     if system_prompt is None and instruction is not None and conclusions is not None and context is not None:
         system_prompt = "Instruction: {}\nConclusions:\n".format(instruction)
         for idx, (conclusion, conclusion_key) in enumerate(conclusions):
-            system_prompt += "{}: {}\n".format(conclusion_key, conclusion)
+            system_prompt += "{}: {}\n".format(conclusion, conclusion_key)
         system_prompt += "\nContext:\n"
         for idx, ctx in enumerate(context):
             system_prompt += "{}: [{}]\n".format(ctx, idx + 1)
@@ -126,8 +127,24 @@ def generate(
     outputs = []
     for text in streamer:
         outputs.append(text)
-        if "[/INST]" in "".join(outputs):
-          return "".join(outputs).replace("[/INST]","")
-        if "[INST]" in "".join(outputs):
-          return "".join(outputs).replace("[INST]","")
-    return "".join(outputs).replace("[INST]","").replace("[/INST]","")
+        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
+        }