Spaces:
Runtime error
Runtime error
| import torch | |
| import gradio as gr | |
| from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
| # 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): | |
| 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 | |
| # Generate code from model and prompt | |
| def generate_code(prompt): | |
| 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 | |
| # 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() |