Update agents/software_engineer_agent.py

agents/software_engineer_agent.py

@@ -5,10 +5,13 @@ MODEL_REPO = "Rahul-8799/software_engineer_mellum"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, torch_dtype=torch.float16, device_map="auto")
-model.eval()
 
-def run(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=512)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def run(state):
+    prompt = state["prompt"]
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    output_ids = model.generate(input_ids, max_new_tokens=512)
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    state["chat_log"].append({"Software Engineer": response})
+    state["pm_out"] = response
+    return state