Update agents/programmer.py

agents/programmer.py

@@ -2,26 +2,24 @@
 
 from agents.base_agent import BaseAgent, ACPMessage
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 
 class ProgrammerAgent(BaseAgent):
     def __init__(self):
         super().__init__(name="CodeBot", role="Programmer")
-        self.tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-        self.model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+        self.tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-multi")
+        self.model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-multi")
 
     def generate_code_reply(self, prompt: str) -> str:
         inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
         outputs = self.model.generate(
             inputs["input_ids"],
-            max_length=100,
+            max_length=128,
             do_sample=True,
-            temperature=0.7,
+            temperature=0.6,
             pad_token_id=self.tokenizer.eos_token_id
         )
         reply = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return reply[len(prompt):].strip()
-
     def receive_message(self, message: ACPMessage) -> ACPMessage:
         if message.performative == "request":
             prompt = f"Write Python code to: {message.content.strip()}\n\n"