fix bug categories

src/api/main.py

@@ -1,13 +1,13 @@
 import os
 import sys
 sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
-
+import json
 import joblib
 import pandas as pd
 import requests
 from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException
-import numpy as np 
+import numpy as np
 from src.data.features import extract_features
 from src.api.schemas import CVERequest, PredictionResponse
 
@@ -18,6 +18,13 @@ api_key = os.getenv("NVD_API_KEY")
 MODEL_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "models", "model_final_xgboost_v2.pkl")
 model = joblib.load(MODEL_PATH)
 
+# The categories created during training are specified
+# Otherwise, the model fails to compare the "sellers" variable and crashes
+CATEGORIES_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "models", "categories.json")
+with open(CATEGORIES_PATH) as f:
+    CATEGORIES = json.load(f)
+
+
 # We define our columns and the Threshold
 CATEGORICAL_COLUMNS = [
     "attack_vector", "attack_complexity", "privileges_required",
@@ -65,8 +72,6 @@ def predict(request: CVERequest):
         raise HTTPException(status_code=502, detail=f"Erreur NVD: {e}")
 
     features = extract_features(cve_raw)
-    # We don't need to include the history
-    # we take the description directly from the time of the CVE's release
     features["description_length_reconstructed"] = features["description_length"]
 
     feature_columns = model.feature_names_in_
@@ -74,8 +79,7 @@ def predict(request: CVERequest):
 
     for col in CATEGORICAL_COLUMNS:
         if col in X.columns:
-            # We use NaN if unknown features
-            X[col] = X[col].where(X[col].notna(), np.nan).astype("category")
+            X[col] = pd.Categorical(X[col], categories=CATEGORIES[col])
 
     proba = model.predict_proba(X)[:, 1][0]