Change model selection

app.py

@@ -1,6 +1,7 @@
 import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
+from huggingface_hub import HfApi
 from threading import Thread
 import gc
 import os
@@ -32,7 +33,6 @@ MODELS = [
     "microsoft/Phi-3-mini-4k-instruct",
     "openai-community/gpt2",
     "openai-community/gpt2-medium",
-    "openai-community/gpt2-large",
     "EleutherAI/pythia-70m",
     "EleutherAI/pythia-160m",
     "EleutherAI/pythia-410m",
@@ -43,9 +43,6 @@ MODELS = [
     "stabilityai/StableLM-Zephyr-3B",
     "NousResearch/Hermes-3-Llama-3.1-8B",
     "meta-llama/Llama-3.2-1B",
-    "meta-llama/Llama-3.2-3B",
-    "THUDM/glm-4-1b-flash",
-    "2Butch/MiniCPM-1B-sft-bf16",
     "SupraLabs/Supra-50M-Base",
     "SupraLabs/Supra-50M-Instruct",
     "SupraLabs/Supra-50M-Reasoning",
@@ -56,8 +53,6 @@ MODELS = [
     "Aravindan/awesome-gpt-2-coder",
     "LiquidAI/LFM2-1.2B",
     "LiquidAI/LFM2-2.6B",
-    "LiquidAI/LFM-350M",
-    "LiquidAI/LFM-700M",
     "LiquidAI/LFM2.5-230M",
     "LiquidAI/LFM2.5-350M",
     "LiquidAI/LFM2.5-1.2B-Instruct",
@@ -106,6 +101,31 @@ class StopOnFlag(StoppingCriteria):
         return model_manager.stop_generation
 
 
+_search_cache = {}
+_last_search_time = 0
+
+def search_hf_models(query):
+    global _last_search_time
+    import time as _time
+    _last_search_time = _time.time()
+
+    if not query or len(query) < 2:
+        return gr.update(choices=MODELS)
+    q = query.strip().lower()
+    if q in _search_cache:
+        return gr.update(choices=_search_cache[q])
+    try:
+        api = HfApi()
+        models = list(api.list_models(search=query, limit=15, sort="downloads"))
+        model_ids = [m.id for m in models if m.id]
+        if model_ids:
+            _search_cache[q] = model_ids
+            return gr.update(choices=model_ids)
+        return gr.update(choices=MODELS)
+    except Exception:
+        return gr.update(choices=MODELS)
+
+
 def get_system_stats(request: gr.Request = None):
     mem = psutil.virtual_memory()
     disk = psutil.disk_usage('/')
@@ -327,10 +347,11 @@ with gr.Blocks(title="SLM Model Tester") as app:
                 gr.Timer(2).tick(get_system_stats, None, stats_output)
 
             with gr.Group():
-                model_id_input = gr.Dropdown(
-                    choices=MODELS, label="Model", allow_custom_value=True,
-                    value=DEFAULT_MODEL
+                model_input = gr.Dropdown(
+                    choices=MODELS, label="Model", value=DEFAULT_MODEL,
+                    allow_custom_value=True,
                 )
+                search_btn = gr.Button("Search HuggingFace", variant="secondary")
                 with gr.Row():
                     load_btn = gr.Button("Load Model", variant="primary", scale=2)
                     clean_btn = gr.Button("Clear Cache", variant="stop", scale=1)
@@ -383,9 +404,15 @@ with gr.Blocks(title="SLM Model Tester") as app:
                 label="Output", lines=15, buttons=["copy"], autoscroll=True
             )
 
+    search_btn.click(
+        fn=search_hf_models,
+        inputs=[model_input],
+        outputs=[model_input],
+    )
+
     load_btn.click(
         fn=load_new_model,
-        inputs=[model_id_input],
+        inputs=[model_input],
         outputs=[status_output]
     )