CodeFixPro / agents /error_fixer.py
ParthBhuptani's picture
Create agents/error_fixer.py
60370a2 verified
raw
history blame
696 Bytes
from transformers import pipeline
import re
# 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()
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:"
try:
result = code_fix_pipeline(prompt)[0]["generated_text"]
return result.strip()
except Exception as e:
return f"Error during fixing: {e}"