import gradio as gr
from transformers import pipeline

# Load model once (important for Spaces)
pipe = pipeline(
    "text-generation",
    model="ibm-granite/granite-3.3-2b-instruct",
    device_map="auto"
)

# Chat function
def generate_response(user_input):
    messages = [
        {"role": "user", "content": user_input},
    ]
    
    output = pipe(messages, max_new_tokens=200)
    
    # Handle output safely
    try:
        response = output[0]["generated_text"][-1]["content"]
    except Exception:
        response = str(output)
    
    return response

# Gradio UI
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Ask something..."),
    outputs=gr.Textbox(label="Response"),
    title="IBM Granite Chatbot",
    description="Powered by granite-3.3-2b-instruct"
)

if __name__ == "__main__":
    interface.launch()