Update app.py

app.py

@@ -1,6 +1,6 @@
-import torch  
+import torch
 import gradio as gr
-from transformers import RobertaConfig, RobertaModel
+from transformers import RobertaConfig, RobertaModel, AutoModelForSeq2SeqLM, AutoTokenizer
 
 # Create a configuration object
 config = RobertaConfig.from_pretrained('roberta-base')
@@ -8,13 +8,13 @@ config = RobertaConfig.from_pretrained('roberta-base')
 # Create the Roberta model
 model = RobertaModel.from_pretrained('roberta-base', config=config)
 
-# Load pretrained model and tokenizer  
+# Load pretrained model and tokenizer
 model_name = "zonghaoyang/DistilRoBERTa-base"
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-#Define function to analyze input code  
-def analyze_code(input_code):                
+# 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()]   
@@ -50,28 +50,28 @@ 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 
+# 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
+# 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)  
+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
+    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