Update app.py

app.py

@@ -1,12 +1,6 @@
 import torch
 import gradio as gr
-from transformers import RobertaConfig, RobertaModel, AutoModelForSeq2SeqLM, AutoTokenizer
-
-# Create a configuration object
-config = RobertaConfig.from_pretrained('roberta-base')
-
-# Create the Roberta model
-model = RobertaModel.from_pretrained('roberta-base', config=config)
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 # Load pretrained model and tokenizer
 model_name = "zonghaoyang/DistilRoBERTa-base"
@@ -15,63 +9,54 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # 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}
+    code_str = " ".join(input_code.split())        
+    sentences = [s.strip() for s in code_str.split(".") if s.strip()]   
+    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 {"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
-	   
+    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 
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    generated_ids = model.generate(input_ids, max_length=100, num_beams=5, early_stopping=True)
+    generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_code
 
 # 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"])
-
-# 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)
-
-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
+    suggestions = ["Use more descriptive variable names", "Add comments to explain complex logic", "Refactor duplicated code into functions"]
+    return suggestions
+
+# Main function to integrate the other functions and generate_code
+def main_function(input_code):
+    code_analysis = analyze_code(input_code)
+    prompt = generate_prompt(code_analysis)
+    generated_code = generate_code(prompt)
+    improvements = suggest_improvements(input_code)
+    return generated_code, improvements
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=main_function,
+    inputs=gr.inputs.Textbox(lines=5, label="Input Code"),
+    outputs=[gr.outputs.Textbox(lines=5, label="Generated Code"), gr.outputs.Textbox(lines=5, label="Suggested Improvements")]
+)
+
+# Launch Gradio interface
+iface.launch()