Spaces:
Runtime error
Runtime error
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) | |
# 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): | |
# 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 | |
# 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 | |