Update app.py

app.py

@@ -2,6 +2,14 @@ import gradio as gr
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from functools import lru_cache
 
+# Pre-selected small models
+MODELS = {
+    "SmolLM2-135M-Instruct": "HuggingFaceTB/SmolLM2-135M-Instruct",
+    "GPT-2 (Small)": "gpt2",
+    "DistilGPT-2": "distilgpt2",
+    "Facebook OPT-125M": "facebook/opt-125m"
+}
+
 # Cache the model and tokenizer to avoid reloading
 @lru_cache(maxsize=1)
 def load_model_cached(model_name):
@@ -13,9 +21,14 @@ def load_model_cached(model_name):
         return f"Error loading model: {str(e)}"
 
 # Function to generate a response from the model
-def chat(model_name, user_input, chat_history, system_prompt=""):
-    if model_name.strip() == "":
-        return "Please enter a valid model name.", chat_history
+def chat(selected_model, user_input, chat_history, system_prompt=""):
+    if not selected_model:
+        return "Please select a model from the dropdown.", chat_history
+
+    # Get the model name from the dropdown
+    model_name = MODELS.get(selected_model)
+    if not model_name:
+        return "Invalid model selected.", chat_history
 
     # Load the model (cached)
     generator = load_model_cached(model_name)
@@ -59,13 +72,13 @@ def chat(model_name, user_input, chat_history, system_prompt=""):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Chat with SmolLM2-135M-Instruct")
+    gr.Markdown("# Chat with Small Language Models")
     
     with gr.Row():
-        model_name = gr.Textbox(
-            label="Enter Hugging Face Model Name",
-            value="HuggingFaceTB/SmolLM2-135M-Instruct",  # Default model
-            placeholder="e.g., HuggingFaceTB/SmolLM2-135M-Instruct"
+        selected_model = gr.Dropdown(
+            label="Select a Model",
+            choices=list(MODELS.keys()),
+            value="SmolLM2-135M-Instruct"  # Default model
         )
     
     chatbot = gr.Chatbot(label="Chat")
@@ -78,7 +91,7 @@ with gr.Blocks() as demo:
     clear_button = gr.Button("Clear Chat")
 
     # Define the chat function
-    user_input.submit(chat, [model_name, user_input, chatbot, system_prompt], [user_input, chatbot])
+    user_input.submit(chat, [selected_model, user_input, chatbot, system_prompt], [user_input, chatbot])
     clear_button.click(lambda: [], None, chatbot, queue=False)
 
 # Launch the app