import gradio as gr
from transformers import pipeline

# Load code-aware model
generator = pipeline("text-generation", model="microsoft/phi-1_5")

# Code assistant function
def generate_from_code(prompt):
    result = generator(
        prompt,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        pad_token_id=generator.tokenizer.eos_token_id
    )[0]["generated_text"]
    
    return result.strip()

# Gradio app with UI and API mode
iface = gr.Interface(
    fn=generate_from_code,
    inputs=gr.Textbox(lines=10, placeholder="Paste code or describe what to do...", label="Prompt / Code"),
    outputs=gr.Textbox(label="AI Output"),
    title="AI Code Helper",
    description="Give a prompt, code snippet, or question. Example: 'Explain this code', 'Complete this function', or 'Fix the bug'."
)

# Enable public API
iface.launch(share=True)