Update src/model_manager.py

src/model_manager.py

@@ -31,52 +31,71 @@ class AutomotiveSLMConfig:
 
 class ModelManager:
     def __init__(self, models_path: str):
+        if not isinstance(models_path, str) or not models_path:
+            raise ValueError(f"models_path must be a non-empty string, got: {models_path!r}")
         self.models_path = models_path
         self.cache = {}
         os.makedirs(self.models_path, exist_ok=True)
 
     def get_available_models(self) -> List[str]:
+        if not os.path.isdir(self.models_path):
+            return []
         files = []
         for f in os.listdir(self.models_path):
+            path = os.path.join(self.models_path, f)
+            if not os.path.isfile(path):
+                continue
             ext = os.path.splitext(f)[1].lower()
             if ext in [".pt", ".pth", ".onnx"]:
                 files.append(f)
         return sorted(files)
 
     def _load_config(self, checkpoint_path: str) -> AutomotiveSLMConfig:
-        # Prefer assets/config.json if present
-        assets_root = os.path.dirname(self.models_path)
+        # Derive assets root safely
+        if not isinstance(checkpoint_path, str):
+            raise ValueError(f"checkpoint_path must be a string, got: {checkpoint_path!r}")
+        assets_root = os.path.dirname(self.models_path)  # assets
         cfg_path = os.path.join(assets_root, "config.json")
-        if os.path.exists(cfg_path):
+        if isinstance(cfg_path, str) and os.path.exists(cfg_path):
             with open(cfg_path, "r") as f:
                 cfg = json.load(f)
             return AutomotiveSLMConfig(**cfg)
-        # else try checkpoint
-        ckpt = torch.load(checkpoint_path, map_location="cpu")
-        if isinstance(ckpt, dict) and "config" in ckpt:
-            return AutomotiveSLMConfig(**ckpt["config"])
+        # Fall back to reading from checkpoint if it’s a torch file
+        ext = os.path.splitext(checkpoint_path)[1].lower()
+        if ext in [".pt", ".pth"] and os.path.exists(checkpoint_path):
+            ckpt = torch.load(checkpoint_path, map_location="cpu")
+            if isinstance(ckpt, dict) and "config" in ckpt:
+                return AutomotiveSLMConfig(**ckpt["config"])
+        # Final fallback
         return AutomotiveSLMConfig()
 
     def load_model(self, model_filename: str) -> Tuple[Any, Any, AutomotiveSLMConfig]:
+        if not isinstance(model_filename, str) or not model_filename:
+            raise ValueError(f"model_filename must be a non-empty string, got: {model_filename!r}")
+
         if model_filename in self.cache:
             return self.cache[model_filename]
+
         model_path = os.path.join(self.models_path, model_filename)
+        if not os.path.isfile(model_path):
+            raise FileNotFoundError(f"Model file not found: {model_path}")
 
-        # tokenizer (GPT-2 per your training)
+        # tokenizer
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
         ext = os.path.splitext(model_filename)[1].lower()
+        config = self._load_config(model_path)
+
         if ext in [".pt", ".pth"]:
-            config = self._load_config(model_path)
             from src.model_architecture import AutomotiveSLM
-            checkpoint = torch.load(model_path, map_location="cpu")
+                        checkpoint = torch.load(model_path, map_location="cpu")
             model = AutomotiveSLM(config)
-            model.load_state_dict(checkpoint["model_state_dict"])
+            state = checkpoint.get("model_state_dict", checkpoint)
+            model.load_state_dict(state, strict=True)
             model.eval()
         elif ext == ".onnx":
-            config = self._load_config(model_path)
             providers = ["CPUExecutionProvider"]
             so = ort.SessionOptions()
             so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
@@ -86,3 +105,4 @@ class ModelManager:
 
         self.cache[model_filename] = (model, tokenizer, config)
         return model, tokenizer, config
+