Upload model_loader.py with huggingface_hub

model_loader.py

@@ -0,0 +1,430 @@
+"""
+Model Loader for the Tiny Conversational AI.
+Handles model downloading, loading, and optimization for CPU inference.
+Uses llama-cpp-python for maximum CPU performance with 4-bit quantization.
+"""
+
+import os
+import sys
+import time
+import shutil
+from pathlib import Path
+from typing import Optional, Dict, Any, Callable
+from dataclasses import dataclass
+import threading
+
+from config import Config, get_config
+from utils import (
+    get_logger, get_system_info, get_available_ram_gb, get_optimal_thread_count,
+    check_system_requirements, ProgressBar, Timer, ensure_dir
+)
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ModelInfo:
+    """Information about a loaded model."""
+    name: str
+    model_id: str
+    path: Path
+    size_gb: float
+    context_size: int
+    loaded: bool = False
+    load_time_seconds: float = 0.0
+
+
+class ModelLoader:
+    """
+    Handles loading and managing LLM models.
+    Uses llama-cpp-python for efficient CPU inference with 4-bit quantization.
+    """
+    
+    def __init__(self, config: Optional[Config] = None):
+        self.config = config or get_config()
+        self.model = None
+        self.model_info: Optional[ModelInfo] = None
+        self._lock = threading.Lock()
+        self._loading = False
+        
+        # Ensure models directory exists
+        ensure_dir(self.config.paths.models_dir)
+    
+    def _get_model_path(self, model_id: str) -> Path:
+        """Get the local path for a model."""
+        model_config = self.config.model.AVAILABLE_MODELS.get(model_id)
+        if not model_config:
+            raise ValueError(f"Unknown model: {model_id}")
+        
+        return self.config.paths.models_dir / model_config["file"]
+    
+    def _is_model_downloaded(self, model_id: str) -> bool:
+        """Check if a model is already downloaded."""
+        path = self._get_model_path(model_id)
+        return path.exists()
+    
+    def _download_model(
+        self,
+        model_id: str,
+        progress_callback: Optional[Callable[[float, str], None]] = None
+    ) -> Path:
+        """
+        Download a model from Hugging Face Hub.
+        Uses huggingface_hub for reliable downloads with resume support.
+        """
+        model_config = self.config.model.AVAILABLE_MODELS.get(model_id)
+        if not model_config:
+            raise ValueError(f"Unknown model: {model_id}")
+        
+        dest_path = self._get_model_path(model_id)
+        
+        if dest_path.exists():
+            logger.info(f"Model already downloaded: {dest_path}")
+            return dest_path
+        
+        logger.info(f"Downloading model: {model_config['name']}")
+        logger.info(f"Repository: {model_config['repo']}")
+        logger.info(f"File: {model_config['file']}")
+        logger.info(f"Expected size: ~{model_config['size_gb']}GB")
+        
+        if progress_callback:
+            progress_callback(0.0, f"Starting download of {model_config['name']}...")
+        
+        try:
+            from huggingface_hub import hf_hub_download
+            
+            # Download with progress
+            downloaded_path = hf_hub_download(
+                repo_id=model_config["repo"],
+                filename=model_config["file"],
+                local_dir=self.config.paths.models_dir,
+                local_dir_use_symlinks=False,
+                resume_download=True,
+            )
+            
+            # Move to expected location if needed
+            downloaded_path = Path(downloaded_path)
+            if downloaded_path != dest_path:
+                if downloaded_path.exists():
+                    shutil.move(str(downloaded_path), str(dest_path))
+            
+            if progress_callback:
+                progress_callback(1.0, "Download complete!")
+            
+            logger.info(f"Model downloaded successfully: {dest_path}")
+            return dest_path
+            
+        except ImportError:
+            logger.error("huggingface_hub not installed. Please install it: pip install huggingface_hub")
+            raise
+        except Exception as e:
+            logger.error(f"Download failed: {e}")
+            raise
+    
+    def select_best_model(self) -> str:
+        """
+        Automatically select the best model based on available RAM.
+        Returns model_id of the selected model.
+        """
+        available_ram = get_available_ram_gb()
+        logger.info(f"Available RAM: {available_ram:.1f}GB")
+        
+        # Sort models by quality (descending), filter by RAM requirement
+        suitable_models = []
+        
+        for model_id, model_config in self.config.model.AVAILABLE_MODELS.items():
+            if model_config["min_ram_gb"] <= available_ram:
+                suitable_models.append((model_id, model_config))
+        
+        if not suitable_models:
+            # Use smallest model as last resort
+            logger.warning("Low RAM detected, using smallest model")
+            return "tinyllama-1.1b"
+        
+        # Sort by quality * speed score
+        suitable_models.sort(
+            key=lambda x: x[1]["quality"] * x[1]["speed"],
+            reverse=True
+        )
+        
+        selected = suitable_models[0][0]
+        logger.info(f"Selected model: {selected}")
+        return selected
+    
+    def load(
+        self,
+        model_id: Optional[str] = None,
+        auto_download: bool = True,
+        progress_callback: Optional[Callable[[float, str], None]] = None
+    ) -> Any:
+        """
+        Load a model for inference.
+        
+        Args:
+            model_id: ID of the model to load. If None, auto-selects best model.
+            auto_download: Whether to download the model if not present.
+            progress_callback: Optional callback for progress updates (progress, message).
+        
+        Returns:
+            The loaded model instance.
+        """
+        with self._lock:
+            if self._loading:
+                raise RuntimeError("Model is already being loaded")
+            self._loading = True
+        
+        try:
+            timer = Timer("Model loading")
+            timer.start()
+            
+            # Auto-select model if not specified
+            if model_id is None:
+                model_id = self.select_best_model()
+            
+            model_config = self.config.model.AVAILABLE_MODELS.get(model_id)
+            if not model_config:
+                raise ValueError(f"Unknown model: {model_id}")
+            
+            model_path = self._get_model_path(model_id)
+            
+            # Download if needed
+            if not model_path.exists():
+                if auto_download:
+                    if progress_callback:
+                        progress_callback(0.0, "Downloading model...")
+                    model_path = self._download_model(model_id, progress_callback)
+                else:
+                    raise FileNotFoundError(
+                        f"Model not found: {model_path}\n"
+                        f"Run with auto_download=True or download manually from: "
+                        f"https://huggingface.co/{model_config['repo']}"
+                    )
+            
+            # Check system requirements (use relaxed check - model will use virtual memory if needed)
+            check_result = check_system_requirements(min(model_config["min_ram_gb"], 1.0))
+            if not check_result["meets_requirements"]:
+                for error in check_result["errors"]:
+                    logger.warning(f"RAM warning (continuing anyway): {error}")
+                # Don't raise - let it try to load, OS will use swap if needed
+            
+            for warning in check_result.get("warnings", []):
+                logger.warning(warning)
+            
+            if progress_callback:
+                progress_callback(0.5, "Loading model into memory...")
+            
+            # Load with llama-cpp-python
+            self.model = self._load_llama_cpp(model_path, model_config)
+            
+            timer.stop()
+            
+            # Store model info
+            self.model_info = ModelInfo(
+                name=model_config["name"],
+                model_id=model_id,
+                path=model_path,
+                size_gb=model_config["size_gb"],
+                context_size=model_config["context_size"],
+                loaded=True,
+                load_time_seconds=timer.elapsed,
+            )
+            
+            if progress_callback:
+                progress_callback(1.0, f"Model loaded in {timer.elapsed:.1f}s")
+            
+            logger.info(f"Model loaded successfully in {timer.elapsed:.1f}s")
+            return self.model
+            
+        finally:
+            with self._lock:
+                self._loading = False
+    
+    def _load_llama_cpp(self, model_path: Path, model_config: Dict[str, Any]) -> Any:
+        """Load model using llama-cpp-python for optimal CPU performance."""
+        try:
+            from llama_cpp import Llama
+        except ImportError:
+            logger.error(
+                "llama-cpp-python not installed. Please install it:\n"
+                "pip install llama-cpp-python"
+            )
+            raise
+        
+        # Determine optimal settings
+        n_threads = self.config.model.n_threads
+        if n_threads == 0:
+            n_threads = get_optimal_thread_count()
+        
+        n_ctx = min(
+            model_config.get("context_size", 4096),
+            self.config.model.max_context_length
+        )
+        
+        logger.info(f"Loading model: {model_path.name}")
+        logger.info(f"Context size: {n_ctx}")
+        logger.info(f"Threads: {n_threads}")
+        logger.info(f"Batch size: {self.config.model.n_batch}")
+        
+        # Load the model
+        model = Llama(
+            model_path=str(model_path),
+            n_ctx=n_ctx,
+            n_threads=n_threads,
+            n_batch=self.config.model.n_batch,
+            n_gpu_layers=self.config.model.n_gpu_layers,
+            use_mmap=self.config.model.use_mmap,
+            use_mlock=self.config.model.use_mlock,
+            verbose=False,  # Reduce noise
+        )
+        
+        return model
+    
+    def unload(self):
+        """Unload the current model to free memory."""
+        with self._lock:
+            if self.model is not None:
+                del self.model
+                self.model = None
+                self.model_info = None
+                
+                # Force garbage collection
+                import gc
+                gc.collect()
+                
+                logger.info("Model unloaded")
+    
+    def is_loaded(self) -> bool:
+        """Check if a model is currently loaded."""
+        return self.model is not None
+    
+    def get_model(self) -> Any:
+        """Get the loaded model instance."""
+        if not self.is_loaded():
+            raise RuntimeError("No model loaded. Call load() first.")
+        return self.model
+    
+    def get_model_info(self) -> Optional[ModelInfo]:
+        """Get information about the loaded model."""
+        return self.model_info
+    
+    def warmup(self, prompt: str = "Hello") -> float:
+        """
+        Warm up the model with a simple generation.
+        Returns the warmup time in seconds.
+        """
+        if not self.is_loaded():
+            raise RuntimeError("No model loaded. Call load() first.")
+        
+        logger.info("Warming up model...")
+        timer = Timer("Warmup")
+        timer.start()
+        
+        # Generate a short response
+        _ = self.model(
+            prompt,
+            max_tokens=10,
+            temperature=0.7,
+        )
+        
+        timer.stop()
+        logger.info(f"Warmup complete in {timer.elapsed:.2f}s")
+        return timer.elapsed
+    
+    def list_available_models(self) -> Dict[str, Dict[str, Any]]:
+        """List all available models with their info."""
+        models = {}
+        
+        for model_id, model_config in self.config.model.AVAILABLE_MODELS.items():
+            models[model_id] = {
+                **model_config,
+                "downloaded": self._is_model_downloaded(model_id),
+                "path": str(self._get_model_path(model_id)),
+            }
+        
+        return models
+    
+    def delete_model(self, model_id: str) -> bool:
+        """Delete a downloaded model to free disk space."""
+        model_path = self._get_model_path(model_id)
+        
+        if model_path.exists():
+            # Don't delete if currently loaded
+            if self.model_info and self.model_info.model_id == model_id:
+                self.unload()
+            
+            model_path.unlink()
+            logger.info(f"Deleted model: {model_path}")
+            return True
+        
+        return False
+
+
+# =============================================================================
+# CONVENIENCE FUNCTIONS
+# =============================================================================
+
+_global_loader: Optional[ModelLoader] = None
+
+
+def get_loader() -> ModelLoader:
+    """Get the global model loader instance."""
+    global _global_loader
+    if _global_loader is None:
+        _global_loader = ModelLoader()
+    return _global_loader
+
+
+def load_model(
+    model_id: Optional[str] = None,
+    auto_download: bool = True
+) -> Any:
+    """Convenience function to load a model."""
+    return get_loader().load(model_id, auto_download)
+
+
+def get_model() -> Any:
+    """Get the currently loaded model."""
+    return get_loader().get_model()
+
+
+if __name__ == "__main__":
+    # Test model loading
+    from utils import print_banner, print_system_status
+    
+    print_banner()
+    print_system_status()
+    
+    loader = ModelLoader()
+    
+    print("\n📦 Available Models:")
+    for model_id, info in loader.list_available_models().items():
+        status = "✓ Downloaded" if info["downloaded"] else "○ Not downloaded"
+        print(f"   • {model_id}: {info['name']}")
+        print(f"     Size: {info['size_gb']}GB | Min RAM: {info['min_ram_gb']}GB | {status}")
+    
+    # Auto-select and load best model
+    print("\n🚀 Loading model...")
+    
+    try:
+        model = loader.load()
+        
+        print(f"\n✓ Model loaded: {loader.model_info.name}")
+        print(f"  Load time: {loader.model_info.load_time_seconds:.1f}s")
+        
+        # Warmup
+        warmup_time = loader.warmup()
+        print(f"  Warmup time: {warmup_time:.2f}s")
+        
+        # Simple test
+        print("\n📝 Test generation:")
+        response = model(
+            "User: Hello!\nAssistant:",
+            max_tokens=50,
+            temperature=0.7,
+            stop=["User:", "\n\n"],
+        )
+        print(f"Response: {response['choices'][0]['text'].strip()}")
+        
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        raise