initial version

.gitignore

@@ -0,0 +1,11 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+.gradio
+
+# Virtual environments
+.venv

.python-version

@@ -0,0 +1 @@
+3.12

main.py

@@ -0,0 +1,96 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, T5ForConditionalGeneration, T5Tokenizer
+
+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 == 'FLAN-T5':
+                self.models['FLAN-T5'] = {
+                    'tokenizer': T5Tokenizer.from_pretrained("google/flan-t5-small"),
+                    'model': T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
+                }
+            
+            # 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 == "FLAN-T5":
+            return self.chat_flan(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_flan(self, message, history):
+        self.ensure_model_loaded('FLAN-T5')
+        
+        tokenizer = self.models['FLAN-T5']['tokenizer']
+        model = self.models['FLAN-T5']['model']
+        
+        inputs = tokenizer(f"Answer the question: {message}", return_tensors="pt")
+        outputs = model.generate(inputs.input_ids, max_length=100)
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+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")
+    
+    with gr.Row():
+        model_dropdown = gr.Dropdown(
+            choices=["SmolLM2", "FLAN-T5"],
+            value="SmolLM2",
+            label="Select Model"
+        )
+    
+    chatbot = gr.Chatbot(height=400)
+    msg = gr.Textbox(label="Message", placeholder="Type your message here...")
+    clear = gr.Button("Clear")
+    
+    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
+    
+    msg.submit(user_message, [msg, chatbot], [msg, chatbot]).then(
+        bot_message, [chatbot, model_dropdown], chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+demo.launch()

pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "small-model-chat"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "gradio>=5.31.0",
+    "huggingface-hub[hf-xet]>=0.31.4",
+    "sentencepiece>=0.2.0",
+    "torch>=2.7.0",
+    "transformers>=4.52.3",
+]