Update src/inference.py

src/inference.py

@@ -11,20 +11,24 @@ sys.path.append(CURRENT_DIR)
 from nlp_utils import BetoClassifier
 from utils import clean_text, preprocess_image_for_ocr
 
+# --- CONFIGURACIÓN ---
 MODEL_VERSION = "v4" 
 MODEL_NAME = "dccuchile/bert-base-spanish-wwm-cased"
 MAX_LEN = 128
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Construcción de la ruta: .../DIMEMEX/models/v4/
-# Subimos dos niveles desde src/inference.py para llegar a la raiz, luego models, luego v4
+# Construcción de la ruta
 PROJECT_ROOT = os.path.abspath(os.path.join(CURRENT_DIR, "../"))
-MODEL_DIR = os.path.join(PROJECT_ROOT, "models", MODEL_VERSION)
+
+# --- CORRECCIÓN AQUÍ ---
+# Antes tenías: os.path.join(PROJECT_ROOT, "models", MODEL_VERSION)
+# Ahora apuntamos directo a 'models', porque ahí viven los archivos .pth
+MODEL_DIR = os.path.join(PROJECT_ROOT, "models")
 
 class MemePredictor:
     def __init__(self):
         print(f"Inicializando motor en: {DEVICE}")
-        print(f"Buscando modelos versión {MODEL_VERSION} en: {MODEL_DIR}")
+        print(f"Buscando modelos en: {MODEL_DIR}")
         
         # Cargar OCR (Singleton)
         self.reader = easyocr.Reader(['es', 'en'], gpu=(DEVICE.type == 'cuda'))
@@ -34,13 +38,13 @@ class MemePredictor:
         self.loaded_models = {}
         
         # Mapas de etiquetas
+        # NOTA: Asegúrate que coincidan con el orden de tu entrenamiento
         self.labels_map = {
             "simple": ["None", "Inappropriate", "Hate"],
-            "complex": ["None", "Inapp", "Sexism", "Racism", "Classicism", "Hate"]
+            "complex": ["None", "Inapp", "Sexism", "Racism", "Classicism", "Other"] 
         }
 
     def _get_model_instance(self, task):
-        # Si ya está en RAM, devolverlo
         if task in self.loaded_models:
             return self.loaded_models[task]
         
@@ -59,7 +63,6 @@ class MemePredictor:
         model = BetoClassifier(n_classes, MODEL_NAME)
         
         # Cargar pesos
-        # map_location es vital para evitar errores si entrenaste en GPU y corres en CPU
         try:
             model.load_state_dict(torch.load(path, map_location=DEVICE))
         except Exception as e: