run.ipynb

{"cells": [{"cell_type": "markdown", "id": "302934307671667531413257853548643485645", "metadata": {}, "source": ["# Gradio Demo: llm_hf_transformers"]}, {"cell_type": "code", "execution_count": null, "id": "272996653310673477252411125948039410165", "metadata": {}, "outputs": [], "source": ["!pip install -q gradio transformers>=4.46.0 torch>=2.3.1 "]}, {"cell_type": "code", "execution_count": null, "id": "288918539441861185822528903084949547379", "metadata": {}, "outputs": [], "source": ["from transformers import AutoModelForCausalLM, AutoTokenizer\n", "import gradio as gr\n", "\n", "checkpoint = \"HuggingFaceTB/SmolLM2-135M-Instruct\"\n", "device = \"cpu\"  # \"cuda\" or \"cpu\"\n", "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n", "model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)\n", "\n", "def predict(message, history):\n", "    history.append({\"role\": \"user\", \"content\": message})\n", "    input_text = tokenizer.apply_chat_template(history, tokenize=False)\n", "    inputs = tokenizer.encode(input_text, return_tensors=\"pt\").to(device)  # type: ignore\n", "    outputs = model.generate(inputs, max_new_tokens=100, temperature=0.2, top_p=0.9, do_sample=True)\n", "    decoded = tokenizer.decode(outputs[0])\n", "    response = decoded.split(\"<|im_start|>assistant\\n\")[-1].split(\"<|im_end|>\")[0]\n", "    return response\n", "\n", "demo = gr.ChatInterface(predict, type=\"messages\")\n", "\n", "if __name__ == \"__main__\":\n", "    demo.launch()\n"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}