app.py

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("""
    <div class="header">
        <h1>🚀 AI Code Generator & Bug Fixer</h1>
        <p>Powered by Advanced AI • Generate, Fix, Optimize & Explain Code</p>
    </div>
    """)
    
    # 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("""
    <div class="footer">
        <h3>🛠️ Built by Hariom Kumar Pandit</h3>
        <p>
            <a href="https://github.com/hari7261" target="_blank">🐙 GitHub: hari7261</a>
            <a href="https://huggingface.co/hari7261" target="_blank">🤗 HuggingFace: hari7261</a>
        </p>
        <p style="margin-top: 1rem; font-size: 0.9rem; opacity: 0.8;">
            Empowering developers with AI-assisted coding • Made with ❤️
        </p>
    </div>
    """)
    
    # 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
    )