import os, time
import gradio as gr
import torch
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer

MODEL_ID = "ibm-granite/granite-4.0-h-tiny"

os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
torch.set_num_threads(2)
torch.set_num_interop_threads(1)

tokenizer = None
model = None

def load_model():
    global tokenizer, model
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
    )
    model.eval()

def respond(message, history, max_tokens):
    messages = list(history or [])
    messages.append({"role": "user", "content": message})

    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt")

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    start = time.perf_counter()
    kwargs = {
        **inputs,
        "streamer": streamer,
        "max_new_tokens": int(max_tokens),
        "do_sample": False,
        "use_cache": True,
    }
    thread = Thread(target=model.generate, kwargs=kwargs)
    thread.start()

    answer = ""
    first_text = None
    for chunk in streamer:
        if first_text is None:
            first_text = time.perf_counter()
        answer += chunk
        yield messages + [{"role": "assistant", "content": answer}], ""

    thread.join()
    elapsed = time.perf_counter() - start
    count = len(tokenizer.encode(answer, add_special_tokens=False))
    speed = count / elapsed if elapsed else 0
    ttft = (first_text - start) * 1000 if first_text else 0

    metrics = (
        f"Output tokens: {count} | Time: {elapsed:.2f}s | "
        f"Approx. speed: {speed:.1f} tok/s | First text: {ttft:.0f} ms"
    )
    yield messages + [{"role": "assistant", "content": answer}], metrics

with gr.Blocks(title="Granite Tiny CPU Lab") as demo:
    gr.Markdown("# Granite 4.0-H-Tiny · CPU Lab")
    chatbot = gr.Chatbot(type="messages", height=440)
    msg = gr.Textbox(label="Message")
    max_tokens = gr.Slider(32, 256, value=128, step=32, label="Max new tokens")
    status = gr.Markdown("Loading model...")

    demo.load(lambda: (load_model(), "Model loaded.")[1], outputs=status)
    msg.submit(
        respond,
        [msg, chatbot, max_tokens],
        [chatbot, status],
    ).then(lambda: "", outputs=msg)

demo.queue(default_concurrency_limit=1).launch(
    server_name="0.0.0.0", server_port=7860
)