Upload alpha_factory/infra/model_manager.py

alpha_factory/infra/model_manager.py

@@ -1,7 +1,7 @@
 """
-Model Manager — Unified interface for Ollama (local) + HuggingFace Inference API (cloud).
-Auto-detects available models from both sources.
-User selects which to use via interactive menu or config.
+Model Manager — Unified interface for Ollama (local + pullable) + HuggingFace Inference API (cloud).
+Auto-detects installed Ollama models AND shows recommended models available to pull.
+User selects which to use via interactive menu, CLI flags, or Gradio dropdowns.
 """
 import asyncio
 import aiohttp
@@ -28,11 +28,34 @@ class ModelInfo:
     quantization: Optional[str] = None
     context_length: Optional[int] = None
     is_default: bool = False
+    is_installed: bool = True  # False = recommended but not yet pulled (Ollama only)
 
     def display_name(self) -> str:
         size_str = f" ({self.size_gb:.1f}GB)" if self.size_gb else ""
         quant_str = f" [{self.quantization}]" if self.quantization else ""
-        return f"[{self.provider.value}] {self.name}{size_str}{quant_str}"
+        pullable = " [PULLABLE — ollama pull " + self.name + "]" if not self.is_installed else ""
+        return f"[{self.provider.value}] {self.name}{size_str}{quant_str}{pullable}"
+
+
+# ─── Ollama models known to work well for this pipeline ────────────────────
+# Includes a range of sizes so every tier has good options.
+OLLAMA_RECOMMENDED = [
+    # Qwen 2.5 family (excellent for structured JSON / codegen)
+    "qwen2.5:0.5b", "qwen2.5:1.5b", "qwen2.5:3b", "qwen2.5:7b",
+    "qwen2.5:14b", "qwen2.5:32b", "qwen2.5:72b",
+    "qwen2.5-coder:1.5b", "qwen2.5-coder:7b", "qwen2.5-coder:14b",
+    # DeepSeek R1 (reasoning-heavy, good for gatekeeper)
+    "deepseek-r1:1.5b", "deepseek-r1:7b", "deepseek-r1:14b",
+    "deepseek-r1:32b", "deepseek-r1:70b",
+    # Llama family
+    "llama3.2:1b", "llama3.2:3b", "llama3.3:70b",
+    # Mistral family
+    "mistral:7b", "mixtral:8x7b", "mixtral:8x22b",
+    # Microsoft Phi
+    "phi4:14b", "phi3:3.8b", "phi3:medium",
+    # Google Gemma
+    "gemma2:2b", "gemma2:9b", "gemma2:27b",
+]
 
 
 # ─── Default model recommendations ─────────────────────────────────────────
@@ -71,6 +94,20 @@ HF_RECOMMENDED = [
 ]
 
 
+# Approximate size mapping for Ollama models (to help tier selection)
+OLLAMA_SIZE_GUESS: dict[str, float] = {
+    "qwen2.5:0.5b": 0.5, "qwen2.5:1.5b": 1.0, "qwen2.5:3b": 2.0,
+    "qwen2.5:7b": 4.7, "qwen2.5:14b": 9.0, "qwen2.5:32b": 20.0, "qwen2.5:72b": 47.0,
+    "qwen2.5-coder:1.5b": 1.0, "qwen2.5-coder:7b": 4.7, "qwen2.5-coder:14b": 9.0,
+    "deepseek-r1:1.5b": 1.0, "deepseek-r1:7b": 4.7, "deepseek-r1:14b": 9.0,
+    "deepseek-r1:32b": 20.0, "deepseek-r1:70b": 43.0,
+    "llama3.2:1b": 0.7, "llama3.2:3b": 2.0, "llama3.3:70b": 43.0,
+    "mistral:7b": 4.7, "mixtral:8x7b": 26.0, "mixtral:8x22b": 80.0,
+    "phi4:14b": 9.0, "phi3:3.8b": 2.5, "phi3:medium": 4.0,
+    "gemma2:2b": 1.6, "gemma2:9b": 5.5, "gemma2:27b": 18.0,
+}
+
+
 def _add_hf_fallbacks(target_list: list[ModelInfo]):
     """Add all HF recommended models as fallbacks."""
     for model_id in HF_RECOMMENDED:
@@ -82,7 +119,9 @@ def _add_hf_fallbacks(target_list: list[ModelInfo]):
 
 class ModelManager:
     """
-    Detects and manages models from Ollama (local) and HuggingFace (cloud).
+    Detects and manages models from:
+    - Ollama (local, installed + recommended-to-pull)
+    - HuggingFace Inference API (cloud)
     Provides unified interface for the pipeline to request models.
     """
 
@@ -110,8 +149,11 @@ class ModelManager:
         )
 
     async def _discover_ollama(self):
-        """Detect locally installed Ollama models."""
+        """Detect locally installed Ollama models AND show recommended models to pull."""
         self.ollama_models = []
+        installed_names: set[str] = set()
+
+        # 1. Query Ollama for already-pulled models
         try:
             async with aiohttp.ClientSession() as session:
                 async with session.get(
@@ -137,8 +179,10 @@ class ModelManager:
                                 provider=ModelProvider.OLLAMA,
                                 size_gb=round(size_gb, 1) if size_gb else None,
                                 quantization=quant,
+                                is_installed=True,
                             ))
-                        logger.info(f"Discovered {len(self.ollama_models)} Ollama models")
+                            installed_names.add(name)
+                        logger.info(f"Discovered {len(self.ollama_models)} installed Ollama models")
                     else:
                         logger.warning(f"Ollama returned status {resp.status}")
         except asyncio.TimeoutError:
@@ -146,6 +190,20 @@ class ModelManager:
         except aiohttp.ClientError as e:
             logger.warning(f"Ollama not reachable: {e}")
 
+        # 2. Add recommended models that are NOT installed (pullable)
+        for tag in OLLAMA_RECOMMENDED:
+            if tag not in installed_names:
+                self.ollama_models.append(ModelInfo(
+                    name=tag,
+                    provider=ModelProvider.OLLAMA,
+                    size_gb=OLLAMA_SIZE_GUESS.get(tag),
+                    is_installed=False,
+                ))
+
+        logger.info(f"Total Ollama choices: {len(self.ollama_models)} "
+                    f"({len(installed_names)} installed + "
+                    f"{len(self.ollama_models) - len(installed_names)} pullable)")
+
     async def _discover_hf(self):
         """Check which HuggingFace models are available via Inference API."""
         self.hf_models = []
@@ -154,7 +212,6 @@ class ModelManager:
             _add_hf_fallbacks(self.hf_models)
             return
 
-        # With token, check which models are actually accessible
         headers = {"Authorization": f"Bearer {self.hf_token}"}
         async with aiohttp.ClientSession() as session:
             for model_id in HF_RECOMMENDED:
@@ -194,13 +251,17 @@ class ModelManager:
                     ))
 
     def get_all_models(self) -> list[ModelInfo]:
-        """Get all discovered models (local + cloud)."""
+        """Get all discovered models (local installed + local pullable + cloud)."""
         return self.ollama_models + self.hf_models
 
     def get_local_models(self) -> list[ModelInfo]:
-        """Get only locally installed models."""
+        """Get Ollama models (installed + pullable)."""
         return self.ollama_models
 
+    def get_installed_models(self) -> list[ModelInfo]:
+        """Get only installed Ollama models."""
+        return [m for m in self.ollama_models if m.is_installed]
+
     def get_cloud_models(self) -> list[ModelInfo]:
         """Get HuggingFace cloud models."""
         return self.hf_models
@@ -227,7 +288,6 @@ class ModelManager:
                 {},
             )
         else:
-            # HuggingFace Inference API
             headers = {}
             if self.hf_token:
                 headers["Authorization"] = f"Bearer {self.hf_token}"
@@ -240,26 +300,33 @@ class ModelManager:
     def auto_assign_defaults(self):
         """
         Automatically assign best available models to each tier.
-        Prefers local (Ollama) over cloud (HF) for speed + privacy.
+        Prefers local installed (Ollama) over pullable over cloud (HF).
         """
-        local_names = {m.name.lower(): m for m in self.ollama_models}
+        installed_names = {m.name.lower(): m for m in self.ollama_models if m.is_installed}
+        all_ollama_names = {m.name.lower(): m for m in self.ollama_models}
 
         for tier, default in DEFAULTS.items():
-            # Try to find the default model locally
-            if default.name.lower() in local_names:
-                self.selected[tier] = local_names[default.name.lower()]
-            elif self.ollama_models:
-                # Use the best available local model for this tier
-                sorted_local = sorted(self.ollama_models, key=lambda m: m.size_gb or 0)
-                if tier == "microfish" and sorted_local:
-                    self.selected[tier] = sorted_local[0]  # smallest
-                elif tier == "bigfish" and sorted_local:
-                    self.selected[tier] = sorted_local[-1]  # largest
-                elif sorted_local:
-                    mid = len(sorted_local) // 2
-                    self.selected[tier] = sorted_local[mid]  # middle
+            # 1. Try exact match among installed
+            if default.name.lower() in installed_names:
+                self.selected[tier] = installed_names[default.name.lower()]
+            # 2. Any installed model, size-appropriate
+            elif installed_names:
+                sorted_installed = sorted(
+                    installed_names.values(),
+                    key=lambda m: m.size_gb or 0
+                )
+                if tier == "microfish" and sorted_installed:
+                    self.selected[tier] = sorted_installed[0]
+                elif tier == "bigfish" and sorted_installed:
+                    self.selected[tier] = sorted_installed[-1]
+                elif sorted_installed:
+                    mid = len(sorted_installed) // 2
+                    self.selected[tier] = sorted_installed[mid]
+            # 3. Fallback to pullable Ollama (same defaults)
+            elif default.name.lower() in all_ollama_names:
+                self.selected[tier] = all_ollama_names[default.name.lower()]
+            # 4. Fallback to HF cloud
             elif self.hf_models:
-                # Fallback to HuggingFace cloud — pick size-appropriate model
                 hf_tier_map = {
                     "microfish": "Qwen/Qwen2.5-7B-Instruct",
                     "tinyfish": "Qwen/Qwen2.5-7B-Instruct",
@@ -270,7 +337,6 @@ class ModelManager:
                 matched = [m for m in self.hf_models if m.name == target]
                 self.selected[tier] = matched[0] if matched else self.hf_models[0]
             else:
-                # Use defaults (will fail at runtime if nothing available)
                 self.selected[tier] = default
 
     def print_status(self):
@@ -284,22 +350,24 @@ class ModelManager:
             console = None
             has_rich = False
 
-        # Discovery summary
+        installed = [m for m in self.ollama_models if m.is_installed]
+        pullable = [m for m in self.ollama_models if not m.is_installed]
+
         if has_rich:
             console.print(f"\n[bold]🔍 Model Discovery[/]")
-            console.print(f"  Ollama (local): {len(self.ollama_models)} models")
+            console.print(f"  Ollama (installed): {len(installed)} models")
+            console.print(f"  Ollama (pullable):  {len(pullable)} models")
             console.print(f"  HuggingFace (cloud): {len(self.hf_models)} models")
             if not self.hf_token:
                 console.print(f"  [yellow]⚠ No HF_TOKEN set — cloud models may have rate limits[/]")
 
-            # Available models table
-            if self.ollama_models:
-                table = Table(title="Local Models (Ollama)")
+            if installed:
+                table = Table(title="Installed Ollama Models")
                 table.add_column("#", width=3)
                 table.add_column("Model", style="cyan")
                 table.add_column("Size", style="green")
                 table.add_column("Quant", style="yellow")
-                for i, m in enumerate(self.ollama_models, 1):
+                for i, m in enumerate(installed, 1):
                     table.add_row(
                         str(i), m.name,
                         f"{m.size_gb:.1f}GB" if m.size_gb else "?",
@@ -307,11 +375,24 @@ class ModelManager:
                     )
                 console.print(table)
 
-            # Selected models
+            if pullable:
+                table = Table(title="Available to Pull (Ollama)")
+                table.add_column("Tag", style="cyan")
+                table.add_column("Est. Size", style="dim")
+                for m in pullable[:15]:  # Limit to avoid wall of text
+                    table.add_row(
+                        m.name,
+                        f"~{m.size_gb:.1f}GB" if m.size_gb else "?",
+                    )
+                if len(pullable) > 15:
+                    table.add_row(f"... and {len(pullable) - 15} more", "")
+                console.print(table)
+
             table2 = Table(title="Selected Models (Pipeline)")
             table2.add_column("Tier", style="bold")
             table2.add_column("Model", style="cyan")
             table2.add_column("Provider", style="magenta")
+            table2.add_column("Status", style="dim")
             table2.add_column("Use", style="dim")
 
             tier_uses = {
@@ -323,18 +404,17 @@ class ModelManager:
 
             for tier in ["microfish", "tinyfish", "mediumfish", "bigfish"]:
                 model = self.get_selected(tier)
+                status = "installed" if model.is_installed else ("pullable" if model.provider == ModelProvider.OLLAMA else "cloud")
                 table2.add_row(
-                    tier, model.name, model.provider.value,
+                    tier, model.name, model.provider.value, status,
                     tier_uses.get(tier, ""),
                 )
             console.print(table2)
         else:
-            # Plain text fallback
             print(f"\nModel Discovery")
-            print(f"  Ollama (local): {len(self.ollama_models)} models")
+            print(f"  Ollama (installed): {len(installed)} models")
+            print(f"  Ollama (pullable):  {len(pullable)} models")
             print(f"  HuggingFace (cloud): {len(self.hf_models)} models")
-            if not self.hf_token:
-                print(f"  ! No HF_TOKEN set — cloud models may have rate limits")
             for tier in ["microfish", "tinyfish", "mediumfish", "bigfish"]:
                 model = self.get_selected(tier)
                 print(f"  {tier}: {model.name} ({model.provider.value})")
@@ -358,7 +438,9 @@ def interactive_model_select(manager: ModelManager) -> dict[str, ModelInfo]:
     all_models = manager.get_all_models()
 
     if not all_models:
-        msg = "No models found! Install Ollama models or set HF_TOKEN.\n  ollama pull qwen2.5:1.5b\n  ollama pull qwen2.5:7b\n  export HF_TOKEN=hf_your_token"
+        msg = ("No models found!\n"
+               "  Install Ollama models: ollama pull qwen2.5:1.5b\n"
+               "  Or set HF_TOKEN for cloud models: export HF_TOKEN=hf_your_token")
         if has_rich:
             console.print(f"[red]{msg}[/]")
         else:
@@ -372,23 +454,26 @@ def interactive_model_select(manager: ModelManager) -> dict[str, ModelInfo]:
     else:
         print("\nAvailable Models:")
         for i, m in enumerate(all_models, 1):
-            print(f"  {i:2d}. [{m.provider.value}] {m.name}")
+            print(f"  {i:2d}. {m.display_name()}")
 
     selections = {}
     for tier in ["microfish", "tinyfish", "mediumfish", "bigfish"]:
         default = DEFAULTS[tier]
-        tier_desc = {"microfish": "bulk generation", "tinyfish": "compilation", "mediumfish": "critique", "bigfish": "final gate"}
+        tier_desc = {
+            "microfish": "bulk generation",
+            "tinyfish": "compilation",
+            "mediumfish": "critique",
+            "bigfish": "final gate",
+        }
         
         if has_rich:
-            console.print(f"\n[bold]Select model for [{tier}][/] (default: {default.name}):")
-            console.print(f"  Use: {tier_desc[tier]}")
+            console.print(f"\n[bold]Select model for [{tier}][/] (default: {default.name}) — {tier_desc[tier]}:")
             choice = Prompt.ask(
                 f"  Enter number (1-{len(all_models)}) or press Enter for default",
                 default="",
             )
         else:
-            print(f"\nSelect model for [{tier}] (default: {default.name}):")
-            print(f"  Use: {tier_desc[tier]}")
+            print(f"\nSelect model for [{tier}] (default: {default.name}) — {tier_desc[tier]}:")
             choice = input(f"  Enter number (1-{len(all_models)}) or press Enter for default: ")
 
         if choice and choice.isdigit():