Updated package v4.1

Author: UtkarshTheDev

locallab/__init__.py

@@ -2,7 +2,7 @@
 LocalLab: Run LLMs locally with a friendly API similar to OpenAI
 """
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 
 from typing import Dict, Any, Optional
 

locallab/model_manager.py

@@ -236,7 +236,12 @@ async def load_model(self, model_id: str) -> bool:
                 if not ENABLE_QUANTIZATION or str(ENABLE_QUANTIZATION).lower() in ('false', '0', 'none', ''):
                     device = "cuda" if torch.cuda.is_available() else "cpu"
                     logger.info(f"Moving model to {device}")
-                    self.model = self.model.to(device)
+                    self.model = AutoModelForCausalLM.from_pretrained(
+                        model_id,
+                        trust_remote_code=True,
+                        token=hf_token,               
+                        device_map="auto"
+                    )
                 else:
                     logger.info("Skipping device move for quantized model - using device_map='auto'")
 

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="locallab",
-    version="0.4.0",
+    version="0.4.1",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=[
         "fastapi>=0.95.0,<1.0.0",