Update app.py

app.py

@@ -1,49 +1,74 @@
 import os
 os.system("pip install -r requirements.txt")
 os.system("pip freeze")
+import torch  
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed, pipeline
 
+# Load pretrained model and tokenizer  
+model_name = "salesforce/codet5-base"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot-small-code-to-text")
-model = AutoModelForCausalLM.from_pretrained("codeparrot/codeparrot-small-code-to-text")
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, num_return_sequences=1, device=-1)
+#Define function to analyze input code  
+def analyze_code(input_code):                
+	# Format code into strings and sentences for NLP     
+	code_str = " ".join(input_code.split())        
+	sentences = [s.strip() for s in code_str.split(".") if s.strip()]   
+	#Extract relevant info and intent from code        
+	variables = []              
+	functions = []    
+	logic = []       
+	for sentence in sentences: 
+		if "=" in sentence:           
+			variables.append(sentence.split("=")[0].strip())       
+		elif "(" in sentence:            
+			functions.append(sentence.split("(")[0].strip())       
+		else:           
+			logic.append(sentence)               
+	#Return info and intent in dictionary    
+	return {"variables": variables, "functions": functions, "logic": logic}
 
+# Define function to generate prompt from analyzed code  
+def generate_prompt(code_analysis):       
+	prompt = f"Generate code with the following: \n\n"   
+	prompt += f"Variables: {', '.join(code_analysis['variables'])} \n\n"   
+	prompt += f"Functions: {', '.join(code_analysis['functions'])} \n\n"   
+	 prompt += f"Logic: {' '.join(code_analysis['logic'])}"  
+	return prompt
+	   
+# Generate code from model and prompt  
+def generate_code(prompt):
+	generated_code = model.generate(prompt, max_length=100, num_beams=5, early_stopping=True)  
+	return generated_code 
 
-def make_doctring(gen_prompt):
-    return gen_prompt + f"\n\n\"\"\"\nExplanation:"
+# Suggest improvements to code
+def suggest_improvements(code):
+	suggestions = ["Use more descriptive variable names", "Add comments to explain complex logic", "Refactor duplicated code into functions"]
+	return suggestions
 
+# Define Gradio interface 
+interface = gr.Interface(fn=generate_code, inputs=["textbox"], outputs=["textbox"])
 
-def code_generation(gen_prompts, max_tokens=8, temperature=0.6, seed=42):
-    set_seed(seed)
-    prompts = [make_doctring(p) for p in gen_prompts]
-    generated_text = pipe(prompts, do_sample=True, top_p=0.95, temperature=temperature, max_length=max_tokens)[0]
-    return generated_text["generated_text"]
+# Have a conversation about the code  
+input_code = """x = 10  
+y = 5   
+def add(a, b):   
+   return a + b
+result = add(x, y)"""
+code_analysis = analyze_code(input_code)        
+prompt = generate_prompt(code_analysis)  
+reply = f"{prompt}\n\n{generate_code(prompt)}\n\nSuggested improvements: {', '.join(suggest_improvements(input_code))}"
+print(reply)
 
-
-title = "Code Explainer"
-description = "This is a space to convert Python code into english text explaining what it does using [codeparrot-small-code-to-text](https://huggingface.co/codeparrot/codeparrot-small-code-to-text),\
-            a code generation model for Python finetuned on [github-jupyter-code-to-text](https://huggingface.co/datasets/codeparrot/github-jupyter-code-to-text) a dataset of Python code followed by a docstring explaining it, the data was originally extracted from Jupyter notebooks."
-
-EXAMPLES = [
-    ["def sort_function(arr):\n    n = len(arr)\n \n    # Traverse through all array elements\n    for i in range(n):\n \n        # Last i elements are already in place\n        for j in range(0, n-i-1):\n \n            # traverse the array from 0 to n-i-1\n            # Swap if the element found is greater\n            # than the next element\n            if arr[j] > arr[j+1]:\n                arr[j], arr[j+1] = arr[j+1], arr[j]"],
-    ["from sklearn import model_selection\nX_train, X_test, Y_train, Y_test = model_selection.train_test_split(X, Y, test_size=0.2)"],
-    ["def load_text(filename):\n    with open(filename, 'r') as f:\n        text = f.read()\n    return text"]
-]
- 
-iface = gr.Interface(
-    fn=code_generation,
-    inputs=[
-        gr.inputs.Textbox(lines=10, placeholder="enter your code here in any programming language...") 
-        gr.inputs.Slider(minimum=8, maximum=256, step=1, default=256, label="Number of tokens to generate"),
-        gr.inputs.Slider(minimum=0, maximum=2.5, step=0.1, default=0.1, label="Temperature"),
-        gr.inputs.Slider(minimum=0, maximum=1000, step=1, default=42, label="Random seed")
-    ],
-    outputs=gr.outputs.Code(language="text", label="Generated explanation", lines=10),
-    examples=EXAMPLES,
-    layout="horizontal",
-    theme="peach",
-    description=description,
-    title=title
-)
-iface.launch()
+while True:
+	change = input("Would you like to make any changes to the code? (Y/N) ")
+	if change == "Y":
+		new_code = input("Enter the updated code: ")
+		code_analysis = analyze_code(new_code)       
+		prompt = generate_prompt(code_analysis)  
+		reply = f"{prompt}\n\n{generate_code(prompt)}\n\nSuggested improvements: {', '.join(suggest_improvements(new_code))}"   
+		print(reply)
+	elif change == "N":
+		print("OK, conversation ended.")
+		break