import gradio as gr
from transformers import pipeline

# Load model
pipe = pipeline(
    "text-generation",
    model="ibm-granite/granite-3b-code-instruct-2k",
    trust_remote_code=True
)

def generate_response(user_input):
    if not user_input.strip():
        return "⚠️ Please enter a prompt."

    # Format messages like your original code
    messages = [
        {"role": "user", "content": user_input}
    ]

    try:
        result = pipe(messages, max_new_tokens=150)
        return result[0]["generated_text"]
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 💻 IBM Granite Code Assistant")

    input_box = gr.Textbox(
        label="Enter your prompt",
        placeholder="Ask something like: Who are you?"
    )

    output_box = gr.Textbox(label="Response")

    btn = gr.Button("Generate")
    btn.click(generate_response, inputs=input_box, outputs=output_box)

# Launch app
if __name__ == "__main__":
    demo.launch()