import gradio as gr
import torch
from transformers import pipeline
import re
import time
# Initialize the model pipeline
model_id = "openai/gpt-oss-120b"
pipe = None
def initialize_model():
global pipe
try:
pipe = pipeline(
"text-generation",
model=model_id,
torch_dtype="auto",
device_map="auto",
)
return "✅ Model loaded successfully!"
except Exception as e:
return f"❌ Error loading model: {str(e)}"
def generate_code(prompt, task_type, language, max_tokens, temperature):
if pipe is None:
return "❌ Model not initialized. Please load the model first.", ""
try:
# Customize prompt based on task type
if task_type == "Generate Code":
system_prompt = f"You are an expert {language} programmer. Generate clean, optimized, and well-commented code for the following request:"
full_prompt = f"{system_prompt}\n\n{prompt}\n\nCode:"
elif task_type == "Fix Bugs":
system_prompt = f"You are an expert {language} debugger. Analyze the following code and fix all bugs, then provide the corrected version:"
full_prompt = f"{system_prompt}\n\n{prompt}\n\nFixed Code:"
elif task_type == "Optimize Code":
system_prompt = f"You are an expert {language} optimizer. Analyze and optimize the following code for better performance and readability:"
full_prompt = f"{system_prompt}\n\n{prompt}\n\nOptimized Code:"
else: # Explain Code
system_prompt = f"You are an expert {language} teacher. Explain the following code step by step:"
full_prompt = f"{system_prompt}\n\n{prompt}\n\nExplanation:"
messages = [
{"role": "user", "content": full_prompt},
]
outputs = pipe(
messages,
max_new_tokens=int(max_tokens),
temperature=temperature,
do_sample=True,
pad_token_id=pipe.tokenizer.eos_token_id
)
generated_text = outputs[0]["generated_text"][-1]["content"] if isinstance(outputs[0]["generated_text"], list) else outputs[0]["generated_text"]
# Extract code if it's wrapped in code blocks
code_match = re.search(r'```(?:\w+\n)?(.*?)```', generated_text, re.DOTALL)
if code_match:
code_output = code_match.group(1).strip()
else:
code_output = generated_text.strip()
# Generate explanation based on the output
explanation = f"Task completed successfully! Generated {len(code_output)} characters of {language} code."
if task_type == "Fix Bugs":
explanation = "Bugs have been identified and fixed. Please review the corrected code."
elif task_type == "Optimize Code":
explanation = "Code has been optimized for better performance and readability."
elif task_type == "Explain Code":
explanation = "Code explanation provided below."
return code_output, explanation
except Exception as e:
return f"❌ Error generating code: {str(e)}", "Please try again with different parameters."
# Custom CSS for modern UI
css = """
.gradio-container {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.header {
text-align: center;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
padding: 2rem;
border-radius: 15px;
margin-bottom: 2rem;
box-shadow: 0 10px 30px rgba(0,0,0,0.2);
}
.header h1 {
font-size: 2.5rem;
font-weight: 700;
margin: 0;
text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
}
.header p {
font-size: 1.2rem;
margin: 0.5rem 0 0 0;
opacity: 0.9;
}
.custom-button {
background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%);
border: none;
color: white;
font-weight: 600;
border-radius: 8px;
transition: all 0.3s ease;
}
.custom-button:hover {
transform: translateY(-2px);
box-shadow: 0 5px 15px rgba(79, 172, 254, 0.4);
}
.footer {
text-align: center;
margin-top: 3rem;
padding: 2rem;
background: linear-gradient(135deg, #2d3436 0%, #636e72 100%);
color: white;
border-radius: 15px;
box-shadow: 0 5px 15px rgba(0,0,0,0.1);
}
.footer h3 {
margin: 0 0 1rem 0;
font-size: 1.3rem;
}
.footer a {
color: #74b9ff;
text-decoration: none;
margin: 0 1rem;
font-weight: 500;
transition: color 0.3s ease;
}
.footer a:hover {
color: #0984e3;
}
.status-box {
padding: 1rem;
border-radius: 8px;
margin: 1rem 0;
font-weight: 500;
}
.code-output {
background: #1e1e1e;
border-radius: 8px;
border: 1px solid #333;
}
.explanation-output {
background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%);
border-radius: 8px;
padding: 1rem;
}
"""
# Create the Gradio interface
with gr.Blocks(css=css, title="AI Code Generator & Bug Fixer", theme=gr.themes.Soft()) as demo:
# Header
gr.HTML("""
""")
# Model initialization section
with gr.Row():
with gr.Column(scale=3):
model_status = gr.Textbox(
label="🤖 Model Status",
value="Click 'Initialize Model' to load the AI model",
interactive=False
)
with gr.Column(scale=1):
init_btn = gr.Button(
"Initialize Model",
variant="primary",
elem_classes=["custom-button"]
)
gr.Markdown("---")
# Main interface
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### ⚙️ Configuration")
task_type = gr.Dropdown(
choices=["Generate Code", "Fix Bugs", "Optimize Code", "Explain Code"],
value="Generate Code",
label="🎯 Task Type"
)
language = gr.Dropdown(
choices=["Python", "JavaScript", "Java", "C++", "C#", "Go", "Rust", "TypeScript", "PHP", "Ruby"],
value="Python",
label="💻 Programming Language"
)
max_tokens = gr.Slider(
minimum=50,
maximum=1000,
value=256,
step=50,
label="📏 Max Tokens"
)
temperature = gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.7,
step=0.1,
label="🌡️ Temperature (Creativity)"
)
generate_btn = gr.Button(
"🚀 Generate/Fix Code",
variant="primary",
size="lg",
elem_classes=["custom-button"]
)
with gr.Column(scale=2):
gr.Markdown("### 📝 Input")
prompt = gr.Textbox(
label="Your Code Request or Buggy Code",
placeholder="Example: Create a function to sort a list of dictionaries by a specific key...",
lines=8
)
gr.Markdown("### 💡 Examples")
examples = gr.Examples(
examples=[
["Generate Code", "Python", "Create a REST API with FastAPI for user management with CRUD operations"],
["Fix Bugs", "JavaScript", "function calculateSum(arr) {\n let sum = 0;\n for (let i = 0; i <= arr.length; i++) {\n sum += arr[i];\n }\n return sum;\n}"],
["Optimize Code", "Python", "def fibonacci(n):\n if n <= 1:\n return n\n return fibonacci(n-1) + fibonacci(n-2)"],
["Explain Code", "Python", "class decorator(func):\n def wrapper(*args, **kwargs):\n print('Before')\n result = func(*args, **kwargs)\n print('After')\n return result\n return wrapper"]
],
inputs=[task_type, language, prompt]
)
gr.Markdown("---")
# Output section
with gr.Row():
with gr.Column():
gr.Markdown("### 📤 Generated Code")
code_output = gr.Code(
label="Result",
language="python",
elem_classes=["code-output"]
)
gr.Markdown("### 💬 AI Explanation")
explanation_output = gr.Textbox(
label="Analysis & Explanation",
lines=3,
elem_classes=["explanation-output"]
)
# Footer
gr.HTML("""
""")
# Event handlers
init_btn.click(
fn=initialize_model,
outputs=model_status
)
generate_btn.click(
fn=generate_code,
inputs=[prompt, task_type, language, max_tokens, temperature],
outputs=[code_output, explanation_output]
)
# Update code language based on selection
def update_code_language(lang):
return gr.Code(language=lang.lower())
language.change(
fn=update_code_language,
inputs=language,
outputs=code_output
)
# Launch the app
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=7860,
share=True,
debug=True
)