Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| class MultiModelChat: | |
| def __init__(self): | |
| self.models = {} | |
| def ensure_model_loaded(self, model_name): | |
| """Lazy load a model only when needed""" | |
| if model_name not in self.models: | |
| print(f"Loading {model_name} model...") | |
| if model_name == 'SmolLM2': | |
| self.models['SmolLM2'] = { | |
| 'tokenizer': AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct"), | |
| 'model': AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct") | |
| } | |
| elif model_name == 'NanoLM-25M': | |
| self.models['NanoLM-25M'] = { | |
| 'tokenizer': AutoTokenizer.from_pretrained("Mxode/NanoLM-25M-Instruct-v1.1"), | |
| 'model': AutoModelForCausalLM.from_pretrained("Mxode/NanoLM-25M-Instruct-v1.1") | |
| } | |
| elif model_name == 'NanoTranslator-S': | |
| self.models['NanoTranslator-S'] = { | |
| 'tokenizer': AutoTokenizer.from_pretrained("Mxode/NanoTranslator-S"), | |
| 'model': AutoModelForCausalLM.from_pretrained("Mxode/NanoTranslator-S") | |
| } | |
| elif model_name == 'NanoTranslator-XL': | |
| self.models['NanoTranslator-XL'] = { | |
| 'tokenizer': AutoTokenizer.from_pretrained("Mxode/NanoTranslator-XL"), | |
| 'model': AutoModelForCausalLM.from_pretrained("Mxode/NanoTranslator-XL") | |
| } | |
| # Set pad token for the newly loaded model | |
| if self.models[model_name]['tokenizer'].pad_token is None: | |
| self.models[model_name]['tokenizer'].pad_token = self.models[model_name]['tokenizer'].eos_token | |
| print(f"{model_name} model loaded successfully!") | |
| def chat(self, message, history, model_choice): | |
| if model_choice == "SmolLM2": | |
| return self.chat_smol(message, history) | |
| elif model_choice == "NanoLM-25M": | |
| return self.chat_nanolm(message, history) | |
| elif model_choice == "NanoTranslator-S": | |
| return self.chat_translator(message, history) | |
| elif model_choice == "NanoTranslator-XL": | |
| return self.chat_translator_xl(message, history) | |
| def chat_smol(self, message, history): | |
| self.ensure_model_loaded('SmolLM2') | |
| tokenizer = self.models['SmolLM2']['tokenizer'] | |
| model = self.models['SmolLM2']['model'] | |
| inputs = tokenizer(f"User: {message}\nAssistant:", return_tensors="pt") | |
| outputs = model.generate( | |
| inputs.input_ids, | |
| max_new_tokens=80, | |
| temperature=0.7, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return response.split("Assistant:")[-1].strip() | |
| def chat_nanolm(self, message, history): | |
| self.ensure_model_loaded('NanoLM-25M') | |
| tokenizer = self.models['NanoLM-25M']['tokenizer'] | |
| model = self.models['NanoLM-25M']['model'] | |
| # Use chat template for NanoLM | |
| messages = [ | |
| {"role": "system", "content": "You are a helpful assistant."}, | |
| {"role": "user", "content": message} | |
| ] | |
| text = tokenizer.apply_chat_template( | |
| messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| inputs = tokenizer([text], return_tensors="pt") | |
| outputs = model.generate( | |
| inputs.input_ids, | |
| max_new_tokens=100, | |
| temperature=0.7, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated_ids = [ | |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, outputs) | |
| ] | |
| response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| return response | |
| def chat_translator(self, message, history): | |
| self.ensure_model_loaded('NanoTranslator-S') | |
| tokenizer = self.models['NanoTranslator-S']['tokenizer'] | |
| model = self.models['NanoTranslator-S']['model'] | |
| # Use translation prompt format | |
| prompt = f"<|im_start|>{message}<|endoftext|>" | |
| inputs = tokenizer([prompt], return_tensors="pt") | |
| outputs = model.generate( | |
| inputs.input_ids, | |
| max_new_tokens=100, | |
| temperature=0.55, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated_ids = [ | |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, outputs) | |
| ] | |
| response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| return response | |
| def chat_translator_xl(self, message, history): | |
| self.ensure_model_loaded('NanoTranslator-XL') | |
| tokenizer = self.models['NanoTranslator-XL']['tokenizer'] | |
| model = self.models['NanoTranslator-XL']['model'] | |
| # Use translation prompt format | |
| prompt = f"<|im_start|>{message}<|endoftext|>" | |
| inputs = tokenizer([prompt], return_tensors="pt") | |
| outputs = model.generate( | |
| inputs.input_ids, | |
| max_new_tokens=100, | |
| temperature=0.55, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated_ids = [ | |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, outputs) | |
| ] | |
| response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| return response | |
| chat_app = MultiModelChat() | |
| def respond(message, history, model_choice): | |
| return chat_app.chat(message, history, model_choice) | |
| with gr.Blocks(theme="soft") as demo: | |
| gr.Markdown("# 🤖 Multi-Model Tiny Chatbot") | |
| gr.Markdown("*Lightweight AI models for different tasks - Choose the right model for your needs!*") | |
| with gr.Row(): | |
| model_dropdown = gr.Dropdown( | |
| choices=["SmolLM2", "NanoLM-25M", "NanoTranslator-S", "NanoTranslator-XL"], | |
| value="NanoLM-25M", | |
| label="Select Model", | |
| info="Choose the best model for your task" | |
| ) | |
| # Model information display | |
| with gr.Row(): | |
| model_info = gr.Markdown( | |
| """ | |
| ## 📋 NanoLM-25M (25M) - Selected | |
| **Best for:** Quick responses, simple tasks, resource-constrained environments | |
| **Language:** English | |
| **Memory:** ~100MB | |
| **Speed:** Very Fast | |
| 💡 **Tip:** Ultra-lightweight model perfect for fast responses! | |
| """, | |
| visible=True | |
| ) | |
| chatbot = gr.Chatbot(height=400, show_label=False) | |
| msg = gr.Textbox( | |
| label="Message", | |
| placeholder="Type your message here...", | |
| lines=2 | |
| ) | |
| with gr.Row(): | |
| clear = gr.Button("🗑️ Clear Chat", variant="secondary") | |
| submit = gr.Button("💬 Send", variant="primary") | |
| # Usage tips | |
| with gr.Accordion("📖 Model Usage Guide", open=False): | |
| gr.Markdown(""" | |
| ### 🎯 When to use each model: | |
| **🔵 SmolLM2 (135M)** | |
| - General conversations and questions | |
| - Creative writing tasks | |
| - Coding help and explanations | |
| - Educational content | |
| **🟢 NanoLM-25M (25M)** | |
| - Quick responses when speed matters | |
| - Resource-constrained environments | |
| - Simple Q&A tasks | |
| - Mobile or edge deployment | |
| **🔴 NanoTranslator-S (9M)** | |
| - Fast English → Chinese translation | |
| - Basic translation needs | |
| - Ultra-low memory usage | |
| - Real-time translation | |
| **🟡 NanoTranslator-XL (78M)** | |
| - High-quality English → Chinese translation | |
| - Professional translation work | |
| - Complex sentences and idioms | |
| - Better context understanding | |
| ### 💡 Pro Tips: | |
| - Models load automatically when first selected (lazy loading) | |
| - Translation models work best with clear, complete sentences | |
| - For translation, input English text and get Chinese output | |
| - Restart the app to free up memory from unused models | |
| """) | |
| def update_model_info(model_choice): | |
| info_map = { | |
| "SmolLM2": """ | |
| ## 📋 SmolLM2 (135M) - Selected | |
| **Best for:** General conversation, Q&A, creative writing, coding help | |
| **Language:** English | |
| **Memory:** ~500MB | |
| **Speed:** Fast | |
| 💡 **Tip:** Great all-around model for most conversational tasks! | |
| """, | |
| "NanoLM-25M": """ | |
| ## 📋 NanoLM-25M (25M) - Selected | |
| **Best for:** Quick responses, simple tasks, resource-constrained environments | |
| **Language:** English | |
| **Memory:** ~100MB | |
| **Speed:** Very Fast | |
| 💡 **Tip:** Ultra-lightweight model perfect for fast responses! | |
| """, | |
| "NanoTranslator-S": """ | |
| ## 📋 NanoTranslator-S (9M) - Selected | |
| **Best for:** Fast English → Chinese translation | |
| **Language:** English → Chinese | |
| **Memory:** ~50MB | |
| **Speed:** Very Fast | |
| 💡 **Tip:** Input English text to get Chinese translation. Great for quick translations! | |
| """, | |
| "NanoTranslator-XL": """ | |
| ## 📋 NanoTranslator-XL (78M) - Selected | |
| **Best for:** High-quality English → Chinese translation | |
| **Language:** English → Chinese | |
| **Memory:** ~300MB | |
| **Speed:** Fast | |
| 💡 **Tip:** Best translation quality for complex sentences and professional use! | |
| """ | |
| } | |
| return info_map.get(model_choice, "") | |
| # Update model info when dropdown changes | |
| model_dropdown.change( | |
| update_model_info, | |
| inputs=[model_dropdown], | |
| outputs=[model_info] | |
| ) | |
| def user_message(message, history): | |
| return "", history + [[message, None]] | |
| def bot_message(history, model_choice): | |
| user_msg = history[-1][0] | |
| bot_response = chat_app.chat(user_msg, history[:-1], model_choice) | |
| history[-1][1] = bot_response | |
| return history | |
| # Handle message submission | |
| msg.submit(user_message, [msg, chatbot], [msg, chatbot]).then( | |
| bot_message, [chatbot, model_dropdown], chatbot | |
| ) | |
| submit.click(user_message, [msg, chatbot], [msg, chatbot]).then( | |
| bot_message, [chatbot, model_dropdown], chatbot | |
| ) | |
| clear.click(lambda: None, None, chatbot, queue=False) | |
| demo.launch() |