ParthBhuptani commited on
Commit
52bd97a
·
verified ·
1 Parent(s): 18e62cf

Create code_modifier.py

Browse files
Files changed (1) hide show
  1. agents/code_modifier.py +17 -0
agents/code_modifier.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
2
+
3
+ # Load the same CodeT5 model
4
+ tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
5
+ model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base")
6
+ modifier_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
7
+
8
+ def modify_code(code: str, instruction: str) -> str:
9
+ """
10
+ Modifies code based on user's instruction prompt.
11
+ """
12
+ prompt = f"Instruction: {instruction}\nCode:\n{code}\nModified Code:"
13
+ try:
14
+ result = modifier_pipeline(prompt, max_length=256)[0]["generated_text"]
15
+ return result.strip()
16
+ except Exception as e:
17
+ return f"Error during modification: {e}"