Update agents/error_fixer.py

agents/error_fixer.py

@@ -1,14 +1,13 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("bigcode/santacoder")
-model = AutoModelForCausalLM.from_pretrained("bigcode/santacoder")
-fix_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-small")
+model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-small")
+fix_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
 def fix_code(code: str) -> str:
-    prompt = f"""# The following Python code has bugs. Fix the code and output only the corrected version.\n\n# Buggy Code:\n{code}\n\n# Fixed Code:\n"""
+    prompt = f"""The following Python code has some bugs. Fix the code:\n\n{code}\n\nCorrected version:"""
     try:
-        result = fix_pipeline(prompt, max_length=256, do_sample=True)[0]["generated_text"]
-        # Post-process to strip out everything before "# Fixed Code"
-        return result.split("# Fixed Code:")[-1].strip()
+        result = fix_pipeline(prompt, max_length=256)[0]["generated_text"]
+        return result.strip()
     except Exception as e:
-        return f"Error during fixing: {e}"
+        return f"Error during fixing: {e}"