ParthBhuptani commited on
Commit
3326c2a
·
verified ·
1 Parent(s): 6b0f6ee

Update agents/error_fixer.py

Browse files
Files changed (1) hide show
  1. agents/error_fixer.py +8 -1
agents/error_fixer.py CHANGED
@@ -6,7 +6,14 @@ model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base")
6
  code_fix_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
7
 
8
  def fix_code(code: str) -> str:
9
- prompt = f"Here is some buggy code:\n{code}\n\nPlease correct the code and return only the fixed version."
 
 
 
 
 
 
 
10
  try:
11
  result = code_fix_pipeline(prompt, max_length=256)[0]["generated_text"]
12
  return result.strip()
 
6
  code_fix_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
7
 
8
  def fix_code(code: str) -> str:
9
+ prompt = f"""You are a code-fixing AI assistant.
10
+ Below is a Python code snippet that contains bugs. Your job is to fix the code and output the corrected version.
11
+
12
+ [BUGGY CODE]
13
+ {code}
14
+
15
+ [FIXED CODE]
16
+ """
17
  try:
18
  result = code_fix_pipeline(prompt, max_length=256)[0]["generated_text"]
19
  return result.strip()