Update coordinator/task_parser.py

coordinator/task_parser.py

@@ -2,13 +2,14 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-MODEL_NAME = "facebook/opt-125m"
+# Use a larger model for better reasoning
+MODEL_NAME = "facebook/opt-1.3b"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 def parse_brief_with_reasoning(brief: str):
     """
-    Use LLM to extract tasks and show reasoning step-by-step
+    Use LLM to extract tasks and reasoning
     Returns a dict:
         {
             "reasoning": <llm reasoning>,
@@ -16,24 +17,35 @@ def parse_brief_with_reasoning(brief: str):
         }
     """
     prompt = f"""
-You are an AI assistant. A user gives you a project brief. 
-Explain your reasoning step-by-step about what technical tasks you would extract from this brief.
-Then provide the final list of tasks.
+You are an expert project manager. Given the project brief, list the distinct technical tasks needed to implement it.
+Explain your reasoning step by step.
+Output format:
+Reasoning:
+- Explain step 1
+- Explain step 2
+Tasks:
+1. First task
+2. Second task
+Do not repeat the project brief.
 
 Project Brief:
 {brief}
-
-Reasoning and Tasks:
 """
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=200)
+    outputs = model.generate(**inputs, max_new_tokens=300)
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Split reasoning and final tasks
+    # Split reasoning and tasks
     if "Tasks:" in decoded:
         reasoning, tasks_text = decoded.split("Tasks:", 1)
     else:
         reasoning, tasks_text = decoded, ""
 
-    tasks = [line.strip("- ").strip() for line in tasks_text.split("\n") if line.strip()]
+    # Deduplicate tasks
+    tasks = []
+    for line in tasks_text.split("\n"):
+        task = line.strip("-123. ").strip()
+        if task and task not in tasks:
+            tasks.append(task)
+
     return {"reasoning": reasoning.strip(), "tasks": tasks}