Update app.py

app.py

@@ -1,80 +1,382 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-import json, os
+import json
+import os
 
-MODEL_FILE = "models.json"
+DEFAULT_FILE = "default_models.json"
+USER_FILE = "models.json"
 
-def load_models():
-    if os.path.exists(MODEL_FILE):
-        return json.load(open(MODEL_FILE))
-    return {"Default (GPT‑OSS‑20B)": "openai/gpt-oss-20b"}
 
-def save_models(models):
-    json.dump(models, open(MODEL_FILE, "w"))
+# -----------------------------
+# Model data loading / saving
+# -----------------------------
+def load_default_models():
+    with open(DEFAULT_FILE, "r", encoding="utf-8") as f:
+        return json.load(f)
 
-models = load_models()
 
+def load_user_models():
+    if os.path.exists(USER_FILE):
+        with open(USER_FILE, "r", encoding="utf-8") as f:
+            try:
+                return json.load(f)
+            except json.JSONDecodeError:
+                return {}
+    return {}
 
-def add_model(link):
-    name = link.split("/")[-1]
-    models[name] = link
-    save_models(models)
-    return gr.Dropdown.update(choices=list(models.keys()), value=name)
 
+def save_user_models(data):
+    with open(USER_FILE, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
 
+
+def merge_models():
+    """
+    Merge default + user models into one tree:
+      Category -> Family -> Model -> meta
+    User models can introduce new categories/families.
+    """
+    base = load_default_models()
+    user = load_user_models()
+
+    for category, families in user.items():
+        if category not in base:
+            base[category] = {}
+        for family, models in families.items():
+            if family not in base[category]:
+                base[category][family] = {}
+            for model_name, meta in models.items():
+                base[category][family][model_name] = meta
+
+    return base
+
+
+# -----------------------------
+# Utility: flatten and lookup
+# -----------------------------
+def flatten_models(model_tree):
+    """
+    Returns a dict:
+      full_key -> (meta, category, family, model_name)
+    where full_key = "Category / Family / Model"
+    """
+    flat = {}
+    for category, families in model_tree.items():
+        for family, models in families.items():
+            for model_name, meta in models.items():
+                full_key = f"{category} / {family} / {model_name}"
+                flat[full_key] = (meta, category, family, model_name)
+    return flat
+
+
+# -----------------------------
+# Add a new model
+# -----------------------------
+def add_model_box(
+    category,
+    family,
+    model_name,
+    model_id,
+    description,
+    link,
+    emoji
+):
+    # Basic validation
+    if not model_id:
+        return gr.Markdown.update(
+            value="Please provide a Model ID like `user/model`."
+        )
+
+    # Fallbacks
+    if not category:
+        category = "Custom"
+    if not family:
+        family = "User-Added"
+    if not model_name:
+        model_name = model_id.split("/")[-1]
+    if not description:
+        description = "User-added model."
+    if not link:
+        link = f"https://huggingface.co/{model_id}"
+    if not emoji:
+        emoji = "✨"
+
+    user_models = load_user_models()
+
+    # Ensure nested structure exists
+    if category not in user_models:
+        user_models[category] = {}
+    if family not in user_models[category]:
+        user_models[category][family] = {}
+
+    user_models[category][family][model_name] = {
+        "id": model_id,
+        "description": description,
+        "link": link,
+        "emoji": emoji
+    }
+
+    save_user_models(user_models)
+
+    msg = (
+        f"Added model under `{category} / {family}`: "
+        f"{emoji} **{model_name}** (`{model_id}`)\n\n"
+        f"It will appear in the model tree after reloading the Space."
+    )
+    return gr.Markdown.update(value=msg)
+
+
+# -----------------------------
+# Chat function
+# -----------------------------
 def respond(
     message,
-    history: list[dict[str, str]],
+    history,
     system_message,
     max_tokens,
     temperature,
     top_p,
-    model_choice,
-    hf_token: gr.OAuthToken,
+    active_model_key,
+    hf_token: gr.OAuthToken
 ):
-    model_id = models[model_choice]
+    if active_model_key is None:
+        yield "No model selected. Please choose a model in the sidebar and click 'Use this model'."
+        return
+
+    models = merge_models()
+    flat = flatten_models(models)
+
+    meta_tuple = flat.get(active_model_key)
+    if meta_tuple is None:
+        yield "Selected model not found. Please choose a model again."
+        return
+
+    meta, _, _, _ = meta_tuple
+    model_id = meta["id"]
+
     client = InferenceClient(token=hf_token.token, model=model_id)
 
     messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)
+    messages.extend(history or [])
     messages.append({"role": "user", "content": message})
 
     response = ""
-
-    for message in client.chat_completion(
+    for msg in client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
         temperature=temperature,
         top_p=top_p,
     ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
+        delta = msg.choices[0].delta.content or ""
+        response += delta
         yield response
 
 
+# -----------------------------
+# Build the sidebar tree
+# -----------------------------
+def build_model_tree(
+    models,
+    active_model_state,
+    current_model_label,
+    info_markdowns,
+    use_buttons
+):
+    """
+    models: merged models dict (Category -> Family -> Model -> meta)
+    active_model_state: gr.State storing current active full key
+    current_model_label: gr.Markdown for 'Current model: ...'
+    info_markdowns: dict full_key -> gr.Markdown (Model Info)
+    use_buttons: dict full_key -> gr.Button (Use this model)
+    """
+
+    flat = flatten_models(models)
+
+    # We’ll build the UI, and inside each loop create small closures
+    for category, families in models.items():
+        with gr.Accordion(category, open=False):
+            for family, model_dict in families.items():
+                with gr.Accordion(family, open=False):
+                    for model_name, meta in model_dict.items():
+                        emoji = meta.get("emoji", "✨")
+                        full_key = f"{category} / {family} / {model_name}"
+
+                        # Model button row
+                        model_button = gr.Button(f"{emoji} {model_name}", size="sm")
+
+                        # Model info area
+                        info_md = gr.Markdown(visible=False)
+                        info_markdowns[full_key] = info_md
+
+                        # Use button
+                        use_btn = gr.Button("Use this model", visible=False, size="sm")
+                        use_buttons[full_key] = use_btn
+
+                        # ---- Click model button → show this info, hide others
+                        def show_info(clicked_key, fk=full_key):
+                            models_local = merge_models()
+                            flat_local = flatten_models(models_local)
+
+                            updates = {}
+                            for key, (meta_loc, _, _, _) in flat_local.items():
+                                md_out = info_markdowns.get(key)
+                                btn_out = use_buttons.get(key)
+                                if md_out is None or btn_out is None:
+                                    continue
+
+                                if key == fk:
+                                    text = (
+                                        f"**Model ID:** `{meta_loc['id']}`  \n"
+                                        f"**Description:** {meta_loc['description']}  \n"
+                                        f"[Model card]({meta_loc['link']})"
+                                    )
+                                    updates[md_out] = gr.Markdown.update(value=text, visible=True)
+                                    updates[btn_out] = gr.Button.update(visible=True)
+                                else:
+                                    updates[md_out] = gr.Markdown.update(visible=False)
+                                    updates[btn_out] = gr.Button.update(visible=False)
+                            return updates
+
+                        model_button.click(
+                            show_info,
+                            inputs=active_model_state,
+                            outputs=list(info_markdowns.values()) + list(use_buttons.values()),
+                        )
+
+                        # ---- Click "Use this model" → set active model + label
+                        def use_model(fk=full_key):
+                            models_local = merge_models()
+                            flat_local = flatten_models(models_local)
+                            meta_loc_tuple = flat_local.get(fk)
+
+                            if not meta_loc_tuple:
+                                return fk, gr.Markdown.update(
+                                    value="**Current model:** _none selected_"
+                                )
+
+                            meta_loc, _, _, mname = meta_loc_tuple
+                            emoji_local = meta_loc.get("emoji", "✨")
+                            label_text = f"**Current model:** {emoji_local} {mname}"
+                            return fk, gr.Markdown.update(value=label_text)
+
+                        use_btn.click(
+                            use_model,
+                            inputs=None,
+                            outputs=[active_model_state, current_model_label],
+                        )
+
+
+# -----------------------------
+# Build the UI
+# -----------------------------
 with gr.Blocks() as demo:
+    models_tree = merge_models()
+
+    # Holds full key: "Category / Family / Model"
+    active_model_key = gr.State(value=None)
+
     with gr.Sidebar():
         gr.LoginButton()
-        gr.Markdown("### Add a new model")
-        new_model = gr.Textbox(label="Model repo (e.g. meta-llama/Llama-3-8B)")
-        model_dropdown = gr.Dropdown(list(models.keys()), label="Choose model")
-        add_button = gr.Button("Add model")
-        add_button.click(add_model, new_model, model_dropdown)
 
+        # Collapsible "Add New Model" box
+        with gr.Accordion("Add New Model", open=False):
+            category_input = gr.Textbox(
+                label="Category (e.g. Exotic or new category)",
+                placeholder="Exotic"
+            )
+            family_input = gr.Textbox(
+                label="Family (e.g. RWKV)",
+                placeholder="RWKV"
+            )
+            model_name_input = gr.Textbox(
+                label="Model Name (e.g. RWKV-World-7B)",
+                placeholder="RWKV-World-7B"
+            )
+            model_id_input = gr.Textbox(
+                label="Model ID (e.g. BlinkDL/rwkv-7-world)",
+                placeholder="BlinkDL/rwkv-7-world"
+            )
+            description_input = gr.Textbox(
+                label="Description (optional)",
+                lines=2
+            )
+            link_input = gr.Textbox(
+                label="Link (optional, will default to https://huggingface.co/ModelID if empty)",
+                lines=1
+            )
+            emoji_input = gr.Textbox(
+                label="Emoji (optional, e.g. 🌍)",
+                lines=1
+            )
+
+            add_button = gr.Button("Add Model")
+            add_status = gr.Markdown("")
+
+            add_button.click(
+                add_model_box,
+                inputs=[
+                    category_input,
+                    family_input,
+                    model_name_input,
+                    model_id_input,
+                    description_input,
+                    link_input,
+                    emoji_input,
+                ],
+                outputs=add_status,
+            )
+
+        # Current model label under the Add box
+        current_model_label = gr.Markdown("**Current model:** _none selected_")
+
+        # Structures to hold info areas
+        info_markdowns = {}
+        use_buttons = {}
+
+        gr.Markdown("### Models")
+
+        # Build nested accordions for models
+        build_model_tree(
+            models_tree,
+            active_model_state=active_model_key,
+            current_model_label=current_model_label,
+            info_markdowns=info_markdowns,
+            use_buttons=use_buttons,
+        )
+
+    # Main chat interface
     chatbot = gr.ChatInterface(
         respond,
+        title="chatbot",
         type="messages",
         additional_inputs=[
-            gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-            gr.Slider(minimum=1, maximum=2048, value=512, 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"),
-            model_dropdown,
+            gr.Textbox(
+                value="You are a friendly chatbot.",
+                label="System message"
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=2048,
+                value=512,
+                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"
+            ),
+            active_model_key,  # passes current active model key into respond()
         ],
     )