Update main.py

main.py

@@ -21,8 +21,8 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 # ============================================================================
 
 MODEL_PATH = os.getenv("MODEL_PATH", "./saved_model")
-TOKENIZER_NAME = "UBC-NLP/MARBERT"   # نحمّل الـ tokenizer مباشرة من Hub
-MAX_LENGTH = 256
+TOKENIZER_NAME = "UBC-NLP/MARBERT"
+MAX_LENGTH = 128   # نفس القيمة في التدريب الأصلي
 
 LABELS_EN = ["Low", "Medium", "High", "Critical"]
 LABELS_AR = ["منخفضة", "متوسطة", "عالية", "حرجة"]
@@ -50,11 +50,23 @@ async def lifespan(app: FastAPI):
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
         model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
+
+        # تشخيص أحجام المفردات
+        tokenizer_vocab = len(tokenizer)
+        model_vocab = model.config.vocab_size
+        print(f"[diagnostic] Tokenizer vocab size: {tokenizer_vocab}")
+        print(f"[diagnostic] Model vocab size:     {model_vocab}")
+
+        # مزامنة الأحجام إذا كان فيه اختلاف
+        if tokenizer_vocab != model_vocab:
+            print(f"[fix] Resizing model token embeddings: {model_vocab} -> {tokenizer_vocab}")
+            model.resize_token_embeddings(tokenizer_vocab)
+
         model.to(device).eval()
         state["tokenizer"] = tokenizer
         state["model"] = model
         state["device"] = device
-        print(f"[startup] Model ready on {device} | num_labels={model.config.num_labels}")
+        print(f"[startup] ✅ Model ready on {device} | num_labels={model.config.num_labels}")
     yield
     state.clear()
 
@@ -92,8 +104,18 @@ def predict_severity(text: str) -> dict:
         raise RuntimeError("Model not loaded.")
 
     device = state["device"]
-    inputs = tokenizer(text, return_tensors="pt", truncation=True,
-                       padding=True, max_length=MAX_LENGTH).to(device)
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=MAX_LENGTH,
+    ).to(device)
+
+    # حماية إضافية: تأكدي من أن جميع الـ token IDs ضمن النطاق
+    vocab_size = model.config.vocab_size
+    inputs["input_ids"] = torch.clamp(inputs["input_ids"], max=vocab_size - 1)
+
     with torch.no_grad():
         logits = model(**inputs).logits
         probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]