Create models.py

models.py

@@ -0,0 +1,28 @@
+# models.py
+
+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():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME, device=device)
+    return 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
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    yi_coder_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16).to(device).eval()
+    return tokenizer, yi_coder_model, device