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 | |