import gradio as gr import torch from functools import lru_cache from transformers import AutoTokenizer, AutoModelForCausalLM try: import spaces except Exception: class _SpacesFallback: @staticmethod def GPU(fn): return fn spaces = _SpacesFallback() MODEL_CHOICES = { "SmolLM2 135M Instruct": "HuggingFaceTB/SmolLM2-135M-Instruct", "SmolLM2 360M Instruct": "HuggingFaceTB/SmolLM2-360M-Instruct", "SmolLM2 1.7B Instruct": "HuggingFaceTB/SmolLM2-1.7B-Instruct", } DEFAULT_MODEL_LABEL = "SmolLM2 360M Instruct" DEFAULT_BACKEND = "GPU" CSS = """ #app-title { font-size: 2.2rem !important; font-weight: 900 !important; margin-bottom: 0.15rem !important; letter-spacing: -0.03em; } #app-subtitle { color: #6b7280; margin-bottom: 1rem; font-size: 1rem; } .control-card { border: 1px solid rgba(128,128,128,0.18); border-radius: 20px; padding: 16px; background: linear-gradient(180deg, rgba(255,255,255,0.06), rgba(255,255,255,0.03)); box-shadow: 0 8px 24px rgba(0,0,0,0.08); } .gradio-container { max-width: 1200px !important; } """ def _get_device(requested: str) -> str: if requested == "cuda" and torch.cuda.is_available(): return "cuda" return "cpu" def _model_device(model) -> torch.device: for param in model.parameters(): if param.device.type != "meta": return param.device return torch.device("cuda" if torch.cuda.is_available() else "cpu") @lru_cache(maxsize=8) def load_model(model_id: str, device: str): device = _get_device(device) tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token if device == "cuda": model = AutoModelForCausalLM.from_pretrained( model_id, dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True, ) else: model = AutoModelForCausalLM.from_pretrained( model_id, dtype=torch.float32, low_cpu_mem_usage=True, ).to(device) model.eval() return tokenizer, model def _extract_text(content) -> str: """Gradio 6's messages format stores `content` as either a plain string or a list of content blocks (e.g. [{"type": "text", "text": "..."}]). Normalize either shape down to plain text.""" if content is None: return "" if isinstance(content, str): return content if isinstance(content, list): parts = [] for block in content: if isinstance(block, str): parts.append(block) elif isinstance(block, dict): if "text" in block: parts.append(block.get("text") or "") return "".join(parts) if isinstance(content, dict): return content.get("text", "") or "" return str(content) def build_messages(history, user_message: str): """history is a list of {'role': ..., 'content': ...} dicts (Gradio Chatbot 'messages' format). content may be a string or a list of content blocks.""" messages = [] for msg in history or []: role = msg.get("role") text = _extract_text(msg.get("content")) if role in ("user", "assistant") and text: messages.append({"role": role, "content": text}) messages.append({"role": "user", "content": user_message}) return messages def build_fallback_prompt(history, user_message: str) -> str: parts = [] for msg in history or []: role = msg.get("role") text = _extract_text(msg.get("content")) if not text: continue if role == "user": parts.append(f"User: {text}") elif role == "assistant": parts.append(f"Assistant: {text}") parts.append(f"User: {user_message}") parts.append("Assistant:") return "\n".join(parts) def _generate(message, history, model_label, max_new_tokens, device_name): model_id = MODEL_CHOICES[model_label] tokenizer, model = load_model(model_id, device_name) max_new_tokens = int(max_new_tokens) if getattr(tokenizer, "chat_template", None): messages = build_messages(history, message) encoded = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, ) input_ids = encoded["input_ids"] else: prompt = build_fallback_prompt(history, message) input_ids = tokenizer(prompt, return_tensors="pt").input_ids input_ids = input_ids.to(_model_device(model)) attention_mask = torch.ones_like(input_ids) generation_kwargs = { "attention_mask": attention_mask, "max_new_tokens": max_new_tokens, "do_sample": True, "temperature": 0.7, "top_p": 0.95, "eos_token_id": tokenizer.eos_token_id, "pad_token_id": tokenizer.eos_token_id, } with torch.inference_mode(): output_ids = model.generate(input_ids, **generation_kwargs) new_tokens = output_ids[0][input_ids.shape[-1]:] reply = tokenizer.decode(new_tokens, skip_special_tokens=True).strip() return reply or "" def generate_cpu(message, history, model_label, max_new_tokens): return _generate(message, history, model_label, max_new_tokens, device_name="cpu") @spaces.GPU def generate_gpu(message, history, model_label, max_new_tokens): return _generate(message, history, model_label, max_new_tokens, device_name="cuda") def respond(message, history, model_label, backend, max_new_tokens): message = (message or "").strip() history = history or [] if not message: return history, "" if backend == "GPU": reply = generate_gpu(message, history, model_label, max_new_tokens) else: reply = generate_cpu(message, history, model_label, max_new_tokens) history = history + [ {"role": "user", "content": message}, {"role": "assistant", "content": reply}, ] return history, "" with gr.Blocks(title="SmolLM2 Chat") as demo: gr.Markdown("# SmolLM2 Chat", elem_id="app-title") gr.Markdown( "A simple chat interface for SmolLM2 instruct models.", elem_id="app-subtitle", ) with gr.Row(): with gr.Column(scale=1, min_width=300): gr.Markdown("### Controls") with gr.Group(elem_classes=["control-card"]): model_dropdown = gr.Dropdown( choices=list(MODEL_CHOICES.keys()), value=DEFAULT_MODEL_LABEL, label="Model", interactive=True, ) backend = gr.Radio( choices=["CPU", "GPU"], value=DEFAULT_BACKEND, label="Backend", interactive=True, ) max_new_tokens = gr.Slider( minimum=32, maximum=2048, value=256, step=1, label="Max new tokens", ) with gr.Column(scale=2, min_width=500): gr.Markdown("### Chat") chatbot = gr.Chatbot( height=620, label="Conversation", ) message = gr.Textbox( placeholder="PauliePocket is just better than you...", label="Message", lines=3, ) gr.Examples( examples=[ ["Explain transformers in ML."], ["Write a story about the 2026 World Cup."], ["Generate a song about a girl falling in love."], ], inputs=message, label="Try these stupid prompts", ) with gr.Row(): send = gr.Button("Send", variant="primary") clear = gr.Button("Clear") send.click( respond, inputs=[message, chatbot, model_dropdown, backend, max_new_tokens], outputs=[chatbot, message], ) message.submit( respond, inputs=[message, chatbot, model_dropdown, backend, max_new_tokens], outputs=[chatbot, message], ) clear.click(lambda: ([], ""), outputs=[chatbot, message]) if __name__ == "__main__": demo.launch(css=CSS, theme=gr.themes.Soft())