Spaces:
Paused
Paused
| 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 | |
| ) |