Upload models/model_manager.py with huggingface_hub

models/model_manager.py

@@ -2,13 +2,19 @@
 
 import os
 import gc
-import torch
 from typing import Dict, List, Tuple, Optional, Any
 from functools import lru_cache
 from pathlib import Path
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-from peft import PeftModel
+# Optional imports for when running with actual models
+try:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import PeftModel
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+    torch = None
 
 from .model_registry import get_model_info, get_all_models, BASE_MODELS
 
@@ -23,6 +29,9 @@ class ModelManager:
         use_4bit: bool = True,
         device_map: str = "auto",
     ):
+        if not TORCH_AVAILABLE:
+            raise ImportError("torch, transformers, peft are required for ModelManager. Install with: pip install torch transformers peft")
+
         self.base_path = Path(base_path) if base_path else Path(__file__).parent.parent.parent
         self.max_cached_models = max_cached_models
         self.use_4bit = use_4bit