rag_engine.py

# ENUNCIADO
# rag_engine.py
# 
# Este archivo contendrá toda la lógica del motor RAG. Se deben
# implementar obligatoriamente las siguientes funciones (con los nombres y parámetros
# exactos que se indican).
# Al inicio del script se deben cargar:
#   • El modelo de embeddings: SentenceTransformer("MongoDB/mdbr-leaf-ir")
#   • El modelo de lenguaje: PleIAs/Pleias-RAG350M (usando AutoTokenizer y AutoModelForCausalLM de transformers).
#   • Los documentos desde documents.json.
# Función recuperar_documentos(consulta, top_k=2, umbral=0.4)
# Dada una consulta en inglés, recupera los documentos más relevantes de la base de conocimiento.
#   • Parámetros:
#       o consulta (str): pregunta del usuario.
#       o top_k (int): número máximo de documentos a retornar.
#       o umbral (float): valor mínimo de similitud (coseno) para considerar un
#         documento relevante. Los documentos con similitud inferior a este
#         umbral se descartan.
#   • Proceso:
#       1. Calcular el embedding de la consulta y de todos los documentos
#          (preferiblemente una sola vez al cargar el script y almacenarlos para
#          evitar recalcular).
#       2. Calcular la similitud del coseno entre el embedding de la consulta y los
#          embeddings de los documentos.
#       3. Ordenar los documentos de mayor a menor similitud.
#       4. Recorrer en ese orden y seleccionar aquellos cuya similitud sea mayor o
#          igual al umbral, hasta un máximo de top_k documentos.
#   • Retorno: Lista con los textos de los documentos seleccionados. 
# Función generar_respuesta(consulta, documentos_recuperados)
# Genera una respuesta usando el modelo de lenguaje, inyectando los documentos
# recuperados como contexto.
# Parámetros:
#   o consulta (str): pregunta original del usuario.
#   o documentos_recuperados (list): lista de textos con los documentos
# relevantes.
# Proceso:
#   1. Se concatenan todos los documentos en un solo string (por ejemplo,
#   separados por espacios).
#   2. Se construye un prompt con el siguiente formato:
#   “””
#   Answer the question based only on the context provided
#   Context: <" ".join(documentos_recuperados)>
#   Question: <consulta>
#   Answer:
#   “””
#   3. Se genera la respuesta con el modelo
# Retorno: Cadena con la respuesta generada.
# Función preguntar(consulta, top_k=2, umbral=0.4)
#   • Descripción:
#       o Función de alto nivel que une la lógica de recuperar_documentos y
#         generar_respuestas
#   • Parámetros: los mismos que recuperar_documentos.
#   • Retorno: La respuesta generada (cadena).