Update agents/error_fixer.py

agents/error_fixer.py

@@ -1,21 +1,14 @@
-from transformers import pipeline
-import re
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load once globally to save time
-code_fix_pipeline = pipeline("text2text-generation", model="bigcode/starcoder", max_length=512)
-
-def clean_code(code):
-    """Basic code cleanup (removes weird chars, extra spaces)."""
-    return re.sub(r'[^\x00-\x7F]+', '', code).strip()
+# 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)
 
 def fix_code(code: str) -> str:
-    """
-    Uses Hugging Face starcoder to attempt fixing the code.
-    """
-    code = clean_code(code)
-    prompt = f"### Fix the following code:\n{code}\n### Fixed code:"
+    prompt = f"Fix the following code:\n{code}"
     try:
-        result = code_fix_pipeline(prompt)[0]["generated_text"]
+        result = code_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}"