dispatchAI/mobile-chat-demo

1import gradio as gr2import spaces3import torch4from transformers import AutoTokenizer, AutoModelForCausalLM5 6MODEL_ID = "dispatchAI/SmolLM2-135M-Instruct-mobile"7 8tokenizer = None9model = None10 11def load_model():12    global tokenizer, model13    if tokenizer is None:14        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)15        model = AutoModelForCausalLM.from_pretrained(16            MODEL_ID,17            torch_dtype=torch.float16,18            device_map="auto",19        )20    return tokenizer, model21 22@spaces.GPU23def chat(message, history):24    tokenizer, model = load_model()25    26    messages = [{"role": "system", "content": "You are a helpful assistant running on a mobile-optimized model."}]27    for h in history:28        messages.append({"role": "user", "content": h[0]})29        messages.append({"role": "assistant", "content": h[1]})30    messages.append({"role": "user", "content": message})31    32    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)33    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)34    35    with torch.no_grad():36        outputs = model.generate(37            **inputs,38            max_new_tokens=256,39            temperature=0.7,40            do_sample=True,41            pad_token_id=tokenizer.eos_token_id,42        )43    44    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)45    return response46 47demo = gr.ChatInterface(48    fn=chat,49    title="🚀 dispatchAI Mobile Chat",50    description="Chat with dispatchAI/SmolLM2-135M-Instruct-mobile — a 135M parameter model optimized for mobile devices. This runs on ZeroGPU.",51    theme=gr.themes.Soft(primary_hue="blue"),52)53 54if __name__ == "__main__":55    demo.launch()56