Update services/model_service.py

services/model_service.py

@@ -21,20 +21,22 @@ class ModelService:
             self._initialized = True
             self._load_models()
 
-    @lru_cache(maxsize=1)
     def _load_models(self):
         try:
             # Load tokenizer
             self.tokenizer = AutoTokenizer.from_pretrained(settings.MODEL_NAME)
-
-            # Load model configuration without modifying rope_scaling
+    
+            # Load model configuration
             config = LlamaConfig.from_pretrained(settings.MODEL_NAME)
-
-            # Remove rope_scaling if present
-            if hasattr(config, "rope_scaling"):
-                logger.info("Removing rope_scaling from configuration...")
-                config.rope_scaling = None
-
+    
+            # Check and update rope_scaling if necessary
+            if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+                logger.info("Updating rope_scaling in configuration...")
+                config.rope_scaling = {
+                    "type": "linear",  # Ensure this matches the expected type
+                    "factor": config.rope_scaling.get('factor', 1.0)  # Use existing factor or default to 1.0
+                }
+    
             # Load model with the updated configuration
             self.model = AutoModelForCausalLM.from_pretrained(
                 settings.MODEL_NAME,
@@ -42,13 +44,14 @@ class ModelService:
                 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