added mps and cput initialization

main/.cache/hub/version.txt

@@ -0,0 +1 @@
+1

main/api.py

@@ -80,48 +80,62 @@ class LLMApi:
             self.logger.error(f"Failed to download model {model_name}: {str(e)}")
             raise
 
-    def initialize_model(self, model_name: str) -> None:
-        """
-        Initialize a model and tokenizer for text generation.
-
-        Args:
-            model_name: The name of the model to initialize
-        """
-        self.logger.info(f"Initializing generation model: {model_name}")
-        try:
-            self.generation_model_name = model_name
-            local_model_path = self.models_path / model_name.split('/')[-1]
-
-            # Check if model exists locally
-            if local_model_path.exists():
-                self.logger.info(f"Loading model from local path: {local_model_path}")
-                model_path = local_model_path
-            else:
-                self.logger.info(f"Loading model from source: {model_name}")
-                model_path = model_name
-
+def initialize_model(self, model_name: str) -> None:
+    """
+    Initialize a model and tokenizer for text generation.
+    Handles different platforms (CUDA, MPS, CPU) appropriately.
+    """
+    self.logger.info(f"Initializing generation model: {model_name}")
+    try:
+        self.generation_model_name = model_name
+        local_model_path = self.models_path / model_name.split('/')[-1]
+
+        # Check if model exists locally
+        if local_model_path.exists():
+            self.logger.info(f"Loading model from local path: {local_model_path}")
+            model_path = local_model_path
+        else:
+            self.logger.info(f"Loading model from source: {model_name}")
+            model_path = model_name
+
+        # Check platform and set appropriate configuration
+        if torch.cuda.is_available():
+            self.logger.info("CUDA detected, using GPU with quantization")
             quantization_config = BitsAndBytesConfig(
                 load_in_8bit=True,
                 llm_int8_threshold=3.0
-                )
-
+            )
             self.generation_model = AutoModelForCausalLM.from_pretrained(
                 model_path,
                 device_map="auto",
                 quantization_config=quantization_config,
                 torch_dtype=torch.float16
-                    )
+            )
+        elif torch.backends.mps.is_available():
+            self.logger.info("Apple Silicon detected, using MPS device")
+            self.generation_model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                device_map="mps",
+                torch_dtype=torch.float16
+            )
+        else:
+            self.logger.info("No GPU detected, falling back to CPU")
+            self.generation_model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                device_map="cpu",
+                torch_dtype=torch.float32  # Use full precision for CPU
+            )
 
-            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-            # Update generation config with tokenizer-specific values
-            self.generation_config["eos_token_id"] = self.tokenizer.eos_token_id
-            self.generation_config["pad_token_id"] = self.tokenizer.eos_token_id
+        # Update generation config with tokenizer-specific values
+        self.generation_config["eos_token_id"] = self.tokenizer.eos_token_id
+        self.generation_config["pad_token_id"] = self.tokenizer.eos_token_id
 
-            self.logger.info(f"Successfully initialized generation model: {model_name}")
-        except Exception as e:
-            self.logger.error(f"Failed to initialize generation model {model_name}: {str(e)}")
-            raise
+        self.logger.info(f"Successfully initialized generation model: {model_name}")
+    except Exception as e:
+        self.logger.error(f"Failed to initialize generation model {model_name}: {str(e)}")
+        raise
 
     def initialize_embedding_model(self, model_name: str) -> None:
         """

main/app.py

@@ -5,8 +5,6 @@ from .routes import router, init_router
 from .utils.logging import setup_logger
 from .utils.validation import validate_hf
 
-
-
 def load_config():
     """Load configuration from yaml file"""
     with open("main/config.yaml", "r") as f: