Update config.py

config.py

@@ -1,36 +1,28 @@
-# config.py
-
-import os
-from pydantic import BaseModel, ConfigDict
-
-# Clase de configuración usando pydantic
-class MyModel(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-# Configuración de la API de Pinecone
-PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
-PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT")
-
-# Configuraciones del modelo
-EMBEDDING_MODEL_NAME = "BAAI/bge-m3"
-EMBED_MODEL = 'BGE_M3-1024'
-DIMENSIONS = 1024
-
-# Nombre del índice de Pinecone
-INDEX_NAME = 'neonatos2'
-
-# Campos de contexto y enlace
-CONTEXT_FIELDS = ['Tag', 'Pregunta', 'Respuesta']
-LINK_FIELDS = ['Tag', 'Respuesta']
-
-# Umbral de similitud por defecto
-SIMILARITY_THRESHOLD_DEFAULT = 0.5
-
-# Prompt del sistema para Yi-Coder
-SYSTEM_PROMPT = """
-Eres un asistente en salud neonatal en Perú, responde y saluda de forma adecuada,
-solo responde en forma de texto de usuario no del chat completo.
-"""
-
-# Longitud máxima por defecto para las respuestas generadas
-MAX_LENGTH_DEFAULT = 100
+# 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