import gradio as gr
from transformers import pipeline
import torch

MODEL_NAME = "ibm-granite/granite-3.3-2b-instruct"

print("Loading model...")

pipe = pipeline(
    "text-generation",
    model=MODEL_NAME,
    device_map="auto",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

print("Model loaded!")

def predict(message, history):

    try:

        outputs = pipe(
            message,
            max_new_tokens=128
        )

        return outputs[0]["generated_text"]

    except Exception as e:
        return f"Error: {str(e)}"


demo = gr.ChatInterface(
    fn=predict,
    title="IBM Granite Chatbot"
)

demo.launch()