Upload alpha_factory/infra/model_manager.py with huggingface_hub

alpha_factory/infra/model_manager.py

@@ -6,10 +6,13 @@ User selects which to use via interactive menu or config.
 import asyncio
 import aiohttp
 import os
+import logging
 from dataclasses import dataclass, field
 from typing import Optional
 from enum import Enum
 
+logger = logging.getLogger(__name__)
+
 
 class ModelProvider(str, Enum):
     OLLAMA = "ollama"
@@ -68,6 +71,15 @@ HF_RECOMMENDED = [
 ]
 
 
+def _add_hf_fallbacks(target_list: list[ModelInfo]):
+    """Add all HF recommended models as fallbacks."""
+    for model_id in HF_RECOMMENDED:
+        target_list.append(ModelInfo(
+            name=model_id,
+            provider=ModelProvider.HUGGINGFACE,
+        ))
+
+
 class ModelManager:
     """
     Detects and manages models from Ollama (local) and HuggingFace (cloud).
@@ -102,7 +114,10 @@ class ModelManager:
         self.ollama_models = []
         try:
             async with aiohttp.ClientSession() as session:
-                async with session.get(f"{self.ollama_url}/api/tags", timeout=aiohttp.ClientTimeout(total=5)) as resp:
+                async with session.get(
+                    f"{self.ollama_url}/api/tags",
+                    timeout=aiohttp.ClientTimeout(total=5)
+                ) as resp:
                     if resp.status == 200:
                         data = await resp.json()
                         for model in data.get("models", []):
@@ -123,51 +138,60 @@ class ModelManager:
                                 size_gb=round(size_gb, 1) if size_gb else None,
                                 quantization=quant,
                             ))
-        except (aiohttp.ClientError, asyncio.TimeoutError):
-            pass  # Ollama not running — that's fine
+                        logger.info(f"Discovered {len(self.ollama_models)} Ollama models")
+                    else:
+                        logger.warning(f"Ollama returned status {resp.status}")
+        except asyncio.TimeoutError:
+            logger.warning("Ollama discovery timed out (5s). Is Ollama running?")
+        except aiohttp.ClientError as e:
+            logger.warning(f"Ollama not reachable: {e}")
 
     async def _discover_hf(self):
         """Check which HuggingFace models are available via Inference API."""
         self.hf_models = []
         if not self.hf_token:
-            # Still list recommended models (user can add token later)
-            for model_id in HF_RECOMMENDED:
-                self.hf_models.append(ModelInfo(
-                    name=model_id,
-                    provider=ModelProvider.HUGGINGFACE,
-                ))
+            logger.info("No HF_TOKEN set — using recommended model list without validation")
+            _add_hf_fallbacks(self.hf_models)
             return
 
         # With token, check which models are actually accessible
-        try:
-            async with aiohttp.ClientSession() as session:
-                headers = {"Authorization": f"Bearer {self.hf_token}"}
-                for model_id in HF_RECOMMENDED:
-                    try:
-                        async with session.get(
-                            f"https://huggingface.co/api/models/{model_id}",
-                            headers=headers,
-                            timeout=aiohttp.ClientTimeout(total=5),
-                        ) as resp:
-                            if resp.status == 200:
-                                data = await resp.json()
-                                self.hf_models.append(ModelInfo(
-                                    name=model_id,
-                                    provider=ModelProvider.HUGGINGFACE,
-                                    context_length=data.get("config", {}).get("max_position_embeddings"),
-                                ))
-                    except:
-                        # Still list it — might work
-                        self.hf_models.append(ModelInfo(
-                            name=model_id,
-                            provider=ModelProvider.HUGGINGFACE,
-                        ))
-        except:
+        headers = {"Authorization": f"Bearer {self.hf_token}"}
+        async with aiohttp.ClientSession() as session:
             for model_id in HF_RECOMMENDED:
-                self.hf_models.append(ModelInfo(
-                    name=model_id,
-                    provider=ModelProvider.HUGGINGFACE,
-                ))
+                try:
+                    async with session.get(
+                        f"https://huggingface.co/api/models/{model_id}",
+                        headers=headers,
+                        timeout=aiohttp.ClientTimeout(total=5),
+                    ) as resp:
+                        if resp.status == 200:
+                            data = await resp.json()
+                            self.hf_models.append(ModelInfo(
+                                name=model_id,
+                                provider=ModelProvider.HUGGINGFACE,
+                                context_length=data.get("config", {}).get("max_position_embeddings"),
+                            ))
+                            logger.debug(f"HF model validated: {model_id}")
+                        elif resp.status == 401:
+                            logger.warning(f"HF token invalid for {model_id}")
+                            self.hf_models.append(ModelInfo(
+                                name=model_id, provider=ModelProvider.HUGGINGFACE
+                            ))
+                        else:
+                            logger.debug(f"HF model {model_id} status {resp.status}")
+                            self.hf_models.append(ModelInfo(
+                                name=model_id, provider=ModelProvider.HUGGINGFACE
+                            ))
+                except asyncio.TimeoutError:
+                    logger.debug(f"HF model {model_id} discovery timed out")
+                    self.hf_models.append(ModelInfo(
+                        name=model_id, provider=ModelProvider.HUGGINGFACE
+                    ))
+                except aiohttp.ClientError as e:
+                    logger.debug(f"HF model {model_id} discovery error: {e}")
+                    self.hf_models.append(ModelInfo(
+                        name=model_id, provider=ModelProvider.HUGGINGFACE
+                    ))
 
     def get_all_models(self) -> list[ModelInfo]:
         """Get all discovered models (local + cloud)."""