Update models.py

models.py

@@ -7,18 +7,23 @@ from config import EMBEDDING_MODEL_NAME
 
 # Cargar el modelo de embeddings
 def load_embedding_model():
-    # CORRECTO: Se determina el dispositivo y se convierte a un objeto torch.device.
+    # Se determina el dispositivo y se convierte a un objeto torch.device.
     device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
     device = torch.device(device_str)
     
-    # Se pasa el objeto device, no la cadena de texto. Esto soluciona el error.
-    embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME, device=device)
+    # LA CORRECCIÓN: Añadir el parámetro 'use_safetensors=True'
+    # Esto fuerza a la librería a cargar 'model.safetensors' en lugar de 'pytorch_model.bin'.
+    embedding_model = SentenceTransformer(
+        EMBEDDING_MODEL_NAME, 
+        device=device, 
+        use_safetensors=True
+    )
+    
     print(f"Embedding model loaded on device: {embedding_model.device}")
     return embedding_model
 
-# Cargar el modelo Yi-Coder
+# Cargar el modelo Yi-Coder (esta función ya estaba bien)
 def load_yi_coder_model():
-    # CORRECTO: Se determina el dispositivo y se convierte a un objeto torch.device.
     device_str = 'cuda' if torch.cuda.is_available() else 'cpu'
     device = torch.device(device_str)
     
@@ -26,13 +31,13 @@ def load_yi_coder_model():
     
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     
-    # El modelo se carga y luego se envía al dispositivo correcto usando .to(device).
     yi_coder_model = AutoModelForCausalLM.from_pretrained(
         model_path,
         torch_dtype=torch.float16,
-        low_cpu_mem_usage=True
+        low_cpu_mem_usage=True,
+        # También es buena práctica añadirlo aquí, aunque no era la causa del error.
+        use_safetensors=True
     ).to(device).eval()
     
     print(f"Yi-Coder model loaded on device: {yi_coder_model.device}")
-    # Se devuelve el objeto device para ser usado en otras partes de la app.
     return tokenizer, yi_coder_model, device