Update models.py

models.py

@@ -4,14 +4,6 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from config import EMBEDDING_MODEL_NAME
-from pydantic import BaseModel
-
-# Clase para los modelos (opcional, si deseas utilizar pydantic)
-class Models(BaseModel):
-    embedding_model: SentenceTransformer
-    tokenizer: AutoTokenizer
-    yi_coder_model: AutoModelForCausalLM
-    device: torch.device
 
 # Cargar el modelo de embeddings
 def load_embedding_model():
@@ -22,7 +14,11 @@ def load_embedding_model():
 # Cargar el modelo Yi-Coder
 def load_yi_coder_model():
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model_path = "01-ai/Yi-Coder-9B-Chat"  # Asegúrate de que esta ruta sea correcta
+    model_path = "01-ai/Yi-Coder-9B-Chat"  # Asegúrate de que esta ruta sea correcta y que el modelo esté disponible
     tokenizer = AutoTokenizer.from_pretrained(model_path)
-    yi_coder_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16).to(device).eval()
+    yi_coder_model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True  # Opcional: ayuda a reducir el uso de memoria al cargar el modelo
+    ).to(device).eval()
     return tokenizer, yi_coder_model, device