Update models.py

models.py

@@ -1,37 +1,45 @@
-# models.py (CON CÓDIGO DE DIAGNÓSTICO)
+# models.py (VERSIÓN A PRUEBA DE ENTORNO ROTO)
 
 import torch
-import sentence_transformers  # Importar la librería completa
+import os
+from huggingface_hub import snapshot_download # Importación clave
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from config import EMBEDDING_MODEL_NAME
 
-# ======================= CÓDIGO DE DIAGNÓSTICO =======================
-# Estas líneas nos dirán la verdad sobre tu entorno.
-print("--- INICIANDO DIAGNÓSTICO DE VERSIONES ---")
-print(f"--- Versión de Sentence-Transformers: {sentence_transformers.__version__}")
-print(f"--- Versión de PyTorch: {torch.__version__}")
-print("------------------------------------------")
-# ====================================================================
-
 # Cargar el modelo de embeddings
 def load_embedding_model():
-    # ... el resto de tu código se mantiene igual ...
+    print("--- Iniciando carga manual del modelo de embedding ---")
+    
+    # 1. Descargar los archivos del modelo a una caché local y obtener la ruta
+    model_folder = snapshot_download(repo_id=EMBEDDING_MODEL_NAME)
+    print(f"Modelo descargado en: {model_folder}")
+
+    # 2. Construir la ruta al archivo problemático
+    problematic_file_path = os.path.join(model_folder, "pytorch_model.bin")
+
+    # 3. Eliminar el archivo .bin si existe, para forzar el uso de .safetensors
+    if os.path.exists(problematic_file_path):
+        print(f"Eliminando archivo problemático: {problematic_file_path}")
+        os.remove(problematic_file_path)
+    else:
+        print("El archivo pytorch_model.bin no existe, se procederá con safetensors.")
+        
+    # 4. Cargar el modelo desde la carpeta local ya "limpia"
+    # Se quita el argumento 'use_safetensors' porque ya no es necesario.
     device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
     device = torch.device(device_str)
     
     embedding_model = SentenceTransformer(
-        EMBEDDING_MODEL_NAME, 
-        device=device, 
-        use_safetensors=True
+        model_folder,  # Cargar desde la ruta local
+        device=device
     )
     
-    print(f"Embedding model loaded on device: {embedding_model.device}")
+    print(f"Modelo de embedding cargado exitosamente desde la ruta local en el dispositivo: {embedding_model.device}")
     return embedding_model
 
-# Cargar el modelo Yi-Coder
+# Cargar el modelo Yi-Coder (se simplifica para consistencia)
 def load_yi_coder_model():
-    # ... el resto de tu código se mantiene igual ...
     device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
     device = torch.device(device_str)
     
@@ -39,12 +47,12 @@ def load_yi_coder_model():
     
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     
+    # Se quita 'use_safetensors' para evitar cualquier posible conflicto.
     yi_coder_model = AutoModelForCausalLM.from_pretrained(
         model_path,
         torch_dtype=torch.float16,
-        low_cpu_mem_usage=True,
-        use_safetensors=True
+        low_cpu_mem_usage=True
     ).to(device).eval()
     
-    print(f"Yi-Coder model loaded on device: {yi_coder_model.device}")
+    print(f"Modelo Yi-Coder cargado en el dispositivo: {yi_coder_model.device}")
     return tokenizer, yi_coder_model, device