ParthBhuptani commited on
Commit
60370a2
·
verified ·
1 Parent(s): 4ffd171

Create agents/error_fixer.py

Browse files
Files changed (1) hide show
  1. agents/error_fixer.py +21 -0
agents/error_fixer.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import re
3
+
4
+ # Load once globally to save time
5
+ code_fix_pipeline = pipeline("text2text-generation", model="bigcode/starcoder", max_length=512)
6
+
7
+ def clean_code(code):
8
+ """Basic code cleanup (removes weird chars, extra spaces)."""
9
+ return re.sub(r'[^\x00-\x7F]+', '', code).strip()
10
+
11
+ def fix_code(code: str) -> str:
12
+ """
13
+ Uses Hugging Face starcoder to attempt fixing the code.
14
+ """
15
+ code = clean_code(code)
16
+ prompt = f"### Fix the following code:\n{code}\n### Fixed code:"
17
+ try:
18
+ result = code_fix_pipeline(prompt)[0]["generated_text"]
19
+ return result.strip()
20
+ except Exception as e:
21
+ return f"Error during fixing: {e}"