Create model_manager.py

services/model_manager.py

@@ -0,0 +1,78 @@
+# model_manager.py
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from llama_cpp import Llama
+from typing import Optional, Dict
+import logging
+from functools import lru_cache
+from config.config import GenerationConfig, ModelConfig
+
+
+class ModelManager:
+    def __init__(self, device: Optional[str] = None):
+        self.logger = logging.getLogger(__name__)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.models: Dict[str, Any] = {}
+        self.tokenizers: Dict[str, Any] = {}
+
+    def load_model(self, model_id: str, model_path: str, model_type: str, config: ModelConfig) -> None:
+        """Load a model with specified configuration."""
+        try:
+            ##could be differnt models, so we can use a factory pattern to load the correct model - textgen, llama, gguf, text2video, text2image etc.
+            if model_type == "llama":
+                self.tokenizers[model_id] = AutoTokenizer.from_pretrained(
+                    model_path,
+                    padding_side='left',
+                    trust_remote_code=True,
+                    **config.tokenizer_kwargs
+                )
+                if self.tokenizers[model_id].pad_token is None:
+                    self.tokenizers[model_id].pad_token = self.tokenizers[model_id].eos_token
+                
+                self.models[model_id] = AutoModelForCausalLM.from_pretrained(
+                    model_path,
+                    device_map="auto",
+                    trust_remote_code=True,
+                    **config.model_kwargs
+                )
+            elif model_type == "gguf":
+                #TODO load the model first from the cache, if not found load the model and save it in the cache
+                #from huggingface_hub import hf_hub_download
+                #prm_model_path = hf_hub_download(
+                #    repo_id="tensorblock/Llama3.1-8B-PRM-Mistral-Data-GGUF",
+                #    filename="Llama3.1-8B-PRM-Mistral-Data-Q4_K_M.gguf"
+                #)
+
+                
+                self.models[model_id] = self._load_quantized_model(
+                    model_path,
+                    **config.quantization_kwargs
+                )
+        except Exception as e:
+            self.logger.error(f"Failed to load model {model_id}: {str(e)}")
+            raise
+
+   
+    def unload_model(self, model_id: str) -> None:
+        """Unload a model and free resources."""
+        if model_id in self.models:
+            del self.models[model_id]
+            if model_id in self.tokenizers:
+                del self.tokenizers[model_id]
+            torch.cuda.empty_cache()
+
+    def _load_quantized_model(self, model_path: str, **kwargs) -> Llama:
+        """Load a quantized GGUF model."""
+        try:
+            n_gpu_layers = -1 if torch.cuda.is_available() else 0
+            model = Llama(
+                model_path=model_path,
+                n_ctx=kwargs.get('n_ctx', 2048),
+                n_batch=kwargs.get('n_batch', 512),
+                n_gpu_layers=kwargs.get('n_gpu_layers', n_gpu_layers),
+                verbose=kwargs.get('verbose', False)
+            )
+            return model
+        except Exception as e:
+            self.logger.error(f"Failed to load GGUF model: {str(e)}")
+            raise