initial commit

.gitignore

@@ -0,0 +1,3 @@
+.venv
+__pycache__
+*.pyc

app.py

@@ -0,0 +1,37 @@
+import gradio as gr
+import rag_engine
+
+def ask(query, top_k, umbral):
+
+    docs = rag_engine.recuperar_documentos(query, top_k=top_k, umbral=umbral)
+    respuesta = rag_engine.generar_respuesta(query, docs)
+
+    docs_formateados = "\n\n---\n\n".join(docs)
+
+    return respuesta, docs_formateados
+
+
+with gr.Blocks() as demo:
+
+    gr.Markdown("# Sistema de preguntas con RAG")
+    gr.Markdown("Haz una pregunta sobre la base de conocimiento.")
+
+    query = gr.Textbox(label="Tu pregunta (en inglés)", placeholder="Where is the hospital?")
+
+    top_k = gr.Slider(1, 5, value=2, step=1, label="Top-k documentos")
+
+    umbral = gr.Slider(0.0, 1.0, value=0.55, step=0.05, label="Umbral de similitud")
+
+    respuesta = gr.Textbox(label="Respuesta", lines=3)
+
+    docs = gr.Textbox(label="Documentos recuperados", lines=6)
+
+    boton = gr.Button("Enviar")
+
+    boton.click(
+        ask,
+        inputs=[query, top_k, umbral],
+        outputs=[respuesta, docs]
+    )
+
+demo.launch()

documents.json

@@ -0,0 +1,7 @@
+{
+    "doc1": "Hospital contact details: You can contact the hospital at email testing@gmail.com, phone +911234567890, or visit us at xyz, abc, 1234, Nepal.",
+    "doc2": "Hospital's working hours: The hospital's working hours are 7:00 AM - 8:00 PM daily.",
+    "doc3": "Official email address: The official email address to contact the hospital is testing@gmail.com.",
+    "doc4": "Main services: We provide comprehensive healthcare services including emergency care, diagnostic testing, surgical procedures, maternity services, and specialized treatments.",
+    "doc5": "Hospital location: The hospital is located at xyz, abc, 1234, Nepal."
+}

rag_engine.py

@@ -0,0 +1,81 @@
+import json
+import torch
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+with open("documents.json", "r", encoding="utf-8") as f:
+    documents = json.load(f)
+
+docs_list = list(documents.values())
+
+embedding_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")
+doc_embeddings = embedding_model.encode(docs_list)
+
+tokenizer = AutoTokenizer.from_pretrained("PleIAs/Pleias-RAG-350M")
+language_model = AutoModelForCausalLM.from_pretrained("PleIAs/Pleias-RAG-350M")
+
+
+def recuperar_documentos(consulta, top_k=2, umbral=0.4):
+    consulta_embedding = embedding_model.encode([consulta])
+    similitudes = cosine_similarity(consulta_embedding, doc_embeddings)[0]
+
+    resultados = []
+
+    for i, score in enumerate(similitudes):
+        if score >= umbral:
+            resultados.append((score, docs_list[i]))
+
+    resultados.sort(key=lambda x: x[0], reverse=True)
+
+    documentos_recuperados = [doc for _, doc in resultados[:top_k]]
+
+    return documentos_recuperados
+
+
+def generar_respuesta(consulta, documentos_recuperados):
+    contexto = " ".join(documentos_recuperados)
+
+    prompt = f"""Answer the question based only on the context provided
+Context: {contexto}
+Question: {consulta}
+Answer:"""
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+
+    with torch.no_grad():
+        output = language_model.generate(
+            **inputs,
+            max_new_tokens=60,
+            do_sample=False,
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    respuesta_completa = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    if "Answer:" in respuesta_completa:
+        respuesta = respuesta_completa.split("Answer:")[-1].strip()
+    else:
+        respuesta = respuesta_completa.strip()
+
+    return respuesta
+
+
+def preguntar(consulta, top_k=2, umbral=0.4):
+    documentos_recuperados = recuperar_documentos(consulta, top_k=top_k, umbral=umbral)
+    respuesta = generar_respuesta(consulta, documentos_recuperados)
+    return respuesta
+
+
+if __name__ == "__main__":
+    pregunta = "Where is the hospital?"
+    docs = recuperar_documentos(pregunta)
+    respuesta = generar_respuesta(pregunta, docs)
+
+    print("Documentos encontrados:")
+    for i, doc in enumerate(docs, start=1):
+        print(f"{i}. {doc}")
+
+    print("\nRespuesta generada:")
+    print(respuesta)

requirements.txt

@@ -0,0 +1,8 @@
+torch
+transformers
+sentence-transformers
+scikit-learn
+fastapi
+uvicorn
+gradio
+pydantic