Update handler.py

handler.py

@@ -74,15 +74,38 @@ class EndpointHandler:
         """Load tokenizer/model from the provided path and attach them to AICore."""
         assert self.ai_core is not None, "AICore must be initialized before loading the model."
 
-        model_id = self.model_path or self.ai_core.model_id or "gpt2"
+        fallback_id = os.getenv("CODETTE_FALLBACK_MODEL_ID", "gpt2")
+        candidate_paths = [
+            self.model_path,
+            os.getenv("CODETTE_MODEL_PATH"),
+            os.getenv("CODETTE_MODEL_ID"),
+            os.path.join("models", "codette-advanced"),
+            os.path.join("models", "codette-v2", "best"),
+        ]
+
+        model_id = (
+            next((c for c in candidate_paths if c and os.path.exists(c)), None)
+            or next((c for c in candidate_paths if c), None)
+            or self.ai_core.model_id
+            or fallback_id
+        )
         self.logger.info("Loading model for AICore from path: %s", model_id)
 
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token or self.tokenizer.unk_token
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token or self.tokenizer.unk_token
 
-        pad_token_id = self.tokenizer.pad_token_id
-        self.model = AutoModelForCausalLM.from_pretrained(model_id, pad_token_id=pad_token_id)
+            pad_token_id = self.tokenizer.pad_token_id
+            self.model = AutoModelForCausalLM.from_pretrained(model_id, pad_token_id=pad_token_id)
+        except Exception as exc:
+            # Fallback to a known-good small model if the provided path is not a model repo
+            self.logger.warning("Model load failed for %s; retrying with fallback %s: %s", model_id, fallback_id, exc)
+            self.tokenizer = AutoTokenizer.from_pretrained(fallback_id)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token or self.tokenizer.unk_token
+            pad_token_id = self.tokenizer.pad_token_id
+            self.model = AutoModelForCausalLM.from_pretrained(fallback_id, pad_token_id=pad_token_id)
 
         if torch and torch.cuda.is_available():
             self.device = "cuda"