| |
| |
| from threading import Thread |
|
|
| import gradio as gr |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer |
|
|
| MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct" |
|
|
| torch.set_num_threads(max(1, torch.get_num_threads())) |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32) |
| model.eval() |
|
|
|
|
| def content_to_text(content): |
| """Gradio 6 message content may be a string or a list of parts; flatten to text.""" |
| if isinstance(content, str): |
| return content |
| if isinstance(content, dict): |
| return content.get("text") or "" |
| if isinstance(content, (list, tuple)): |
| return "".join(content_to_text(p) for p in content) |
| return "" if content is None else str(content) |
|
|
|
|
| def respond( |
| message, |
| history, |
| system_message, |
| max_tokens, |
| temperature, |
| top_p, |
| ): |
| messages = [] |
| if system_message: |
| messages.append({"role": "system", "content": system_message}) |
|
|
| for msg in history or []: |
| if isinstance(msg, dict) and msg.get("role") in ("user", "assistant"): |
| text = content_to_text(msg.get("content")) |
| if text: |
| messages.append({"role": msg["role"], "content": text}) |
|
|
| messages.append({"role": "user", "content": content_to_text(message)}) |
|
|
| inputs = tokenizer.apply_chat_template( |
| messages, add_generation_prompt=True, return_tensors="pt", return_dict=True |
| ) |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) |
| generate_kwargs = dict( |
| **inputs, |
| streamer=streamer, |
| max_new_tokens=int(max_tokens), |
| do_sample=temperature > 0, |
| temperature=float(temperature), |
| top_p=float(top_p), |
| ) |
| thread = Thread(target=model.generate, kwargs=generate_kwargs) |
| thread.start() |
|
|
| response = "" |
| for token in streamer: |
| response += token |
| yield response |
| thread.join() |
|
|
| chatbot = gr.Chatbot(height=600) |
|
|
| demo = gr.ChatInterface( |
| respond, |
| additional_inputs=[ |
| gr.Textbox(value="", label="System message"), |
| gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="Max new tokens"), |
| gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), |
| gr.Slider( |
| minimum=0.1, |
| maximum=1.0, |
| value=0.95, |
| step=0.05, |
| label="Top-P", |
| ), |
|
|
| ], |
| fill_height=True, |
| chatbot=chatbot, |
| ) |
| if __name__ == "__main__": |
| demo.launch(theme="Nymbo/Nymbo_Theme") |
|
|
