Upload alpha_factory/ui.py

alpha_factory/ui.py

@@ -8,7 +8,6 @@ Run: uv run python -m alpha_factory.ui
 import os
 import sys
 import subprocess
-import json
 import asyncio
 import duckdb
 import gradio as gr
@@ -29,6 +28,7 @@ DB_PATH = Path("factor_store/alphas.duckdb")
 # ── Globals (shared across Gradio sessions) ──────────────────────────────────
 _LAST_DISCOVERED_MODELS: list[ModelInfo] = []
 
+
 def _model_choice_name(m: ModelInfo) -> str:
     """Human-readable label for a model in the dropdown."""
     size = f" ({m.size_gb:.1f}GB)" if m.size_gb else ""
@@ -43,7 +43,6 @@ def _discover_models_sync(
     """Synchronous wrapper around async model discovery."""
     global _LAST_DISCOVERED_MODELS
 
-    # Resolve HF token
     token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN", "")
 
     manager = ModelManager(ollama_url=ollama_url, hf_token=token)
@@ -57,12 +56,28 @@ def _discover_models_sync(
 
 
 def _get_dropdown_choices(models: list[ModelInfo]) -> list[str]:
-    """Build dropdown choices: [Use Default (auto-assign)] + discovered models."""
+    """Build dropdown choices: [Use Default] + discovered models."""
     choices = ["Use Default (auto-assign)"]
     choices.extend([_model_choice_name(m) for m in models])
     return choices
 
 
+def refresh_all_dropdowns(ollama_url: str, hf_token: str):
+    """Discover models and return status + 4 identical choice lists."""
+    models = _discover_models_sync(ollama_url=ollama_url, hf_token=hf_token)
+
+    if not models:
+        msg = "⚠️ No models found. Is Ollama running? Is HF_TOKEN set?"
+        return msg, ["Use Default (auto-assign)"], ["Use Default (auto-assign)"], ["Use Default (auto-assign)"], ["Use Default (auto-assign)"]
+
+    local_count = sum(1 for m in models if m.provider == ModelProvider.OLLAMA)
+    cloud_count = sum(1 for m in models if m.provider == ModelProvider.HUGGINGFACE)
+    msg = f"✅ Found {local_count} Ollama + {cloud_count} HuggingFace models"
+
+    choices = _get_dropdown_choices(models)
+    return msg, choices, choices, choices, choices
+
+
 # ── DB helpers (unchanged) ──────────────────────────────────────────────────
 
 def get_alphas_from_db(limit=50):
@@ -117,7 +132,19 @@ def get_full_expression(evt: gr.SelectData):
     return ""
 
 
-# ── Pipeline runner (now with per-tier model overrides) ───────────────────────
+# ── Pipeline runner (with per-tier model overrides) ───────────────────────────
+
+def _extract_model_name(choice: str) -> Optional[str]:
+    """Strip [PROVIDER] prefix and size/quant suffix to get raw model name."""
+    if not choice or choice == "Use Default (auto-assign)":
+        return None
+    if "]" in choice:
+        raw = choice.split("]", 1)[1].strip()
+        # Remove trailing size/quant suffixes: " (4.7GB) [q4_k_m]" etc.
+        raw = raw.split(" (")[0].split(" [")[0].strip()
+        return raw
+    return choice.strip()
+
 
 def _run_pipeline_subprocess(
     batch_size: int,
@@ -136,7 +163,6 @@ def _run_pipeline_subprocess(
     env["NO_COLOR"] = "1"
     env["TERM"] = "dumb"
 
-    # Build CLI args
     cmd = [
         sys.executable, "-m", "alpha_factory.run",
         "--batch-size", str(int(batch_size)),
@@ -148,15 +174,6 @@ def _run_pipeline_subprocess(
     if enable_brain:
         cmd.append("--enable-brain")
 
-    # Only pass per-tier overrides if user selected something other than default
-    def _extract_model_name(choice: str) -> Optional[str]:
-        if not choice or choice == "Use Default (auto-assign)":
-            return None
-        # Strip the [PROVIDER] prefix and size/quant suffix
-        if "]" in choice:
-            return choice.split("]", 1)[1].strip().split(" (")[0].split(" [")[0].strip()
-        return choice.strip()
-
     mf = _extract_model_name(microfish)
     tf = _extract_model_name(tinyfish)
     mmf = _extract_model_name(mediumfish)
@@ -171,8 +188,7 @@ def _run_pipeline_subprocess(
     if bf:
         cmd.extend(["--bigfish", bf])
 
-    # Log the command for debugging
-    print(f"Running: {' '.join(cmd)}")
+    print(f"[UI] Running: {' '.join(cmd)}")
 
     try:
         result = subprocess.run(
@@ -217,23 +233,6 @@ def generate_and_refresh(
     return table, log
 
 
-# ── Model discovery refresh ──────────────────────────────────────────────────
-
-def refresh_model_list(ollama_url: str, hf_token: str) -> tuple[str, list[str]]:
-    """Discover models and return (status_msg, dropdown_choices)."""
-    models = _discover_models_sync(ollama_url=ollama_url, hf_token=hf_token)
-
-    if not models:
-        return "No models found. Is Ollama running? Is HF_TOKEN set?", ["Use Default (auto-assign)"]
-
-    local_count = sum(1 for m in models if m.provider == ModelProvider.OLLAMA)
-    cloud_count = sum(1 for m in models if m.provider == ModelProvider.HUGGINGFACE)
-
-    msg = f"Found {local_count} Ollama + {cloud_count} HF models"
-    choices = _get_dropdown_choices(models)
-    return msg, choices
-
-
 # ── UI Builder ──────────────────────────────────────────────────────────────
 
 def build_ui():
@@ -251,11 +250,13 @@ def build_ui():
                     ollama_url_input = gr.Textbox(
                         value="http://localhost:11434",
                         label="Ollama URL",
+                        info="URL of your local Ollama instance",
                     )
                     hf_token_input = gr.Textbox(
                         value=os.getenv("HF_TOKEN", ""),
                         label="HF Token (optional)",
                         type="password",
+                        info="HuggingFace token for cloud model access",
                     )
                     refresh_models_btn = gr.Button("🔍 Refresh Model List", variant="secondary")
                     discovery_status = gr.Textbox(
@@ -267,63 +268,43 @@ def build_ui():
                 with gr.Column(scale=2):
                     gr.Markdown("### Model Selection — One Per Tier")
                     gr.Markdown("""
-                    | Tier | Role | Typical Size |
-                    |------|------|-------------|
-                    | **Microfish** | Hypothesis generation (bulk) | 1.5B-3B |
-                    | **Tinyfish** | Expression compilation | 3B-7B |
-                    | **Mediumfish** | Crowd scout + Performance surgeon | 7B-14B |
-                    | **Bigfish** | Gatekeeper (final memo) | 14B-72B |
+                    | Tier | Role | Typical Size | Suggested |
+                    |------|------|-------------|-----------|
+                    | **Microfish** | Hypothesis generation (bulk) | 1.5B–3B | `qwen2.5:1.5b` |
+                    | **Tinyfish** | Expression compilation | 3B–7B | `qwen2.5:3b` |
+                    | **Mediumfish** | Crowd scout + Performance surgeon | 7B–14B | `qwen2.5:7b` |
+                    | **Bigfish** | Gatekeeper (final memo) | 14B–72B | `qwen2.5:14b` |
                     """)
 
-                    # Initial choices: just default until discovery
                     default_choices = ["Use Default (auto-assign)"]
 
                     microfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
-                        label="Microfish — Hypothesis Generation",
+                        label="🐟 Microfish — Hypothesis Generation",
+                        info="Bulk idea generation (fast, small model)",
                     )
                     tinyfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
-                        label="Tinyfish — Expression Compilation",
+                        label="🐠 Tinyfish — Expression Compilation",
+                        info="Converts hypothesis to BRAIN expression",
                     )
                     mediumfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
-                        label="Mediumfish — Critique & Diagnosis",
+                        label="🦈 Mediumfish — Critique & Diagnosis",
+                        info="Novelty check + performance analysis",
                     )
                     bigfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
-                        label="Bigfish — Final Gatekeeper",
+                        label="🐋 Bigfish — Final Gatekeeper",
+                        info="Production go/no-go decision (slowest, most capable)",
                     )
 
-            # When refresh is clicked, update all dropdowns
-            refresh_models_btn.click(
-                fn=refresh_model_list,
-                inputs=[ollama_url_input, hf_token_input],
-                outputs=[discovery_status, microfish_dropdown],
-            ).then(
-                lambda choices: gr.Dropdown(choices=choices),
-                inputs=microfish_dropdown,
-                outputs=tinyfish_dropdown,
-            ).then(
-                lambda choices: gr.Dropdown(choices=choices),
-                inputs=microfish_dropdown,
-                outputs=mediumfish_dropdown,
-            ).then(
-                lambda choices: gr.Dropdown(choices=choices),
-                inputs=microfish_dropdown,
-                outputs=bigfish_dropdown,
-            )
-            # Actually the proper way: refresh returns one choices list,
-            # then update all 4 dropdowns with that same list
-            def _update_all_dropdowns(status, choices):
-                return status, choices, choices, choices, choices
-
             refresh_models_btn.click(
-                fn=lambda url, token: _update_all_dropdowns(*refresh_model_list(url, token)),
+                fn=refresh_all_dropdowns,
                 inputs=[ollama_url_input, hf_token_input],
                 outputs=[
                     discovery_status,
@@ -342,13 +323,17 @@ def build_ui():
                         value=3, label="Batch Size", minimum=1, maximum=20,
                     )
                     proven_mode_cb = gr.Checkbox(
-                        value=False, label="Proven Templates (no LLM)",
+                        value=False,
+                        label="Proven Templates (no LLM)",
+                        info="Deterministic generation, guaranteed valid expressions",
                     )
                     enable_brain_cb = gr.Checkbox(
-                        value=False, label="Enable BRAIN Submission (needs token)",
+                        value=False,
+                        label="Enable BRAIN Submission",
+                        info="Requires BRAIN_SESSION_TOKEN env var",
                     )
                     gr.Markdown("---")
-                    gr.Markdown("*Selected models carry over from the Settings tab*")
+                    gr.Markdown("*Model selections carry over from the Settings tab*")
                     generate_btn = gr.Button("Generate New Batch", variant="primary")
                     refresh_table_btn = gr.Button("Refresh Table Only")
 
@@ -417,7 +402,7 @@ def build_ui():
             - Requires `BRAIN_SESSION_TOKEN` from browser devtools
             - Enable "BRAIN Submission" checkbox (disabled by default for safety)
 
-            [GitHub / HuggingFace](https://huggingface.co/gaurv007/alpha-factory)
+            [Repository](https://huggingface.co/gaurv007/alpha-factory)
             """)
 
     return app