Update agents/error_fixer.py

agents/error_fixer.py

@@ -1,21 +1,14 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# Load CodeT5 model (public and free)
-tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
-model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base")
-code_fix_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+tokenizer = AutoTokenizer.from_pretrained("bigcode/santacoder")
+model = AutoModelForCausalLM.from_pretrained("bigcode/santacoder")
+fix_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 def fix_code(code: str) -> str:
-    prompt = f"""You are a code-fixing AI assistant.
-Below is a Python code snippet that contains bugs. Your job is to fix the code and output the corrected version.
-
-[BUGGY CODE]
-{code}
-
-[FIXED CODE]
-"""
+    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"""
     try:
-        result = code_fix_pipeline(prompt, max_length=256)[0]["generated_text"]
-        return result.strip()
+        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()
     except Exception as e:
         return f"Error during fixing: {e}"