Update services/model_service.py

services/model_service.py

@@ -25,32 +25,33 @@ class ModelService:
     @lru_cache(maxsize=1)
     def _load_models(self):
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(settings.MODEL_NAME )
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(settings.MODEL_NAME)
+
+            # Load model configuration and modify rope_scaling if applicable
+            config = LlamaConfig.from_pretrained(settings.MODEL_NAME)
 
-                  # Modify the model configuration to use a valid rope_scaling format
-            config = LlamaConfig.from_pretrained(settings.model_name)
             if hasattr(config, "rope_scaling"):
+                logger.info("Updating rope_scaling configuration...")
                 config.rope_scaling = {
-                    "type": "linear",
-                    "factor": 32.0
+                    "type": "linear",  # Ensure the type is valid
+                    "factor": 32.0    # Ensure factor is a valid float
                 }
-            
-            # Load model with updated configuration
-            #self.model = AutoModelForCausalLM.from_pretrained(model_name, config=config).to(device)
 
-        
-            
+            # Load model with the updated configuration
             self.model = AutoModelForCausalLM.from_pretrained(
                 settings.MODEL_NAME,
                 torch_dtype=torch.float16 if settings.DEVICE == "cuda" else torch.float32,
                 device_map="auto" if settings.DEVICE == "cuda" else None,
                 config=config
             )
+
+            # Load sentence embedder
             self.embedder = SentenceTransformer(settings.EMBEDDER_MODEL)
+
         except Exception as e:
             logger.error(f"Error loading models: {e}")
             raise
 
     def get_models(self):
         return self.tokenizer, self.model, self.embedder
-