import gradio as gr
from transformers import pipeline

# Load model (you can change this to your own)
generator = pipeline("text2text-generation", model="google/flan-t5-base")

# Function to generate code from prompt
def generate_code(prompt):
    result = generator(prompt, max_new_tokens=100)[0]["generated_text"]
    return result.strip()

# Gradio UI and API Interface
iface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(label="Enter prompt (e.g., Describe code)"),
    outputs=gr.Textbox(label="Generated Code"),
    title="Code Generator AI",
    description="This app generates code based on your prompt. You can use it from the browser or call it via API."
)

# Launch Space (UI + API mode enabled)
iface.launch()