|
import gradio as gr |
|
from langchain_community.llms import HuggingFaceHub |
|
from langchain_core.output_parsers import StrOutputParser |
|
from langchain import hub |
|
from langchain_community.document_loaders import PyPDFLoader |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain_community.vectorstores import Chroma |
|
from langchain_huggingface import HuggingFaceEmbeddings |
|
from huggingface_hub import InferenceClient |
|
from rerankers import Reranker |
|
import os |
|
|
|
embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") |
|
persist_directory = "db" |
|
|
|
if not os.path.exists(persist_directory): |
|
loader = PyPDFLoader("Constitucion_española.pdf") |
|
documents = loader.load() |
|
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100) |
|
docs_split = text_splitter.split_documents(documents) |
|
|
|
vectordb = Chroma.from_documents(docs_split, embedding_function, persist_directory=persist_directory) |
|
vectordb.persist() |
|
else: |
|
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function) |
|
|
|
client = InferenceClient("mrm8488/spanish-t5-small-sqac-for-qa", token=os.getenv("HUGGINGFACEHUB_API_TOKEN")) |
|
|
|
ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type='colbert') |
|
|
|
def generate_text(context, query): |
|
prompt = f"""Responde en español. Basándote en el siguiente contexto contesta de forma clara y directa. |
|
|
|
Contexto: |
|
{context} |
|
|
|
Pregunta: {query} |
|
Respuesta:""" |
|
response = client.text_generation(prompt=prompt, max_new_tokens=200) |
|
return response.strip() |
|
|
|
def test_rag_reranking(query, ranker): |
|
print(f"\n🔍 Pregunta recibida: {query}") |
|
|
|
docs = vectordb.similarity_search_with_score(query) |
|
print(f"🔎 Documentos recuperados: {len(docs)}") |
|
|
|
context = [] |
|
for i, (doc, score) in enumerate(docs): |
|
print(f"📄 Doc {i} - Score: {score}") |
|
if score < 7: |
|
doc_details = doc.to_json()['kwargs'] |
|
content = doc_details['page_content'] |
|
context.append(content) |
|
print(f"✅ Doc {i} agregado al contexto") |
|
|
|
if not context: |
|
print("❌ No se encontró contexto relevante.") |
|
return "No se encontró información suficiente para responder." |
|
|
|
print(f"📚 Contextos pasados al ranker: {len(context)}") |
|
|
|
context_strings = [str(c) for c in context] |
|
|
|
reranked = ranker.rank(query=query, docs=context_strings) |
|
|
|
print(f"🏅 Resultado del reranker: {reranked}") |
|
|
|
best_context = reranked[0].document.text |
|
print(f"🧠 Contexto elegido: {best_context[:500]}...") |
|
|
|
respuesta = generate_text(best_context, query) |
|
print(f"💬 Respuesta generada: {respuesta}") |
|
|
|
return respuesta |
|
|
|
def responder_chat(message, history): |
|
respuesta = test_rag_reranking(message, ranker) |
|
return respuesta |
|
|
|
demo = gr.ChatInterface( |
|
fn=responder_chat, |
|
title="Chatbot sobre la constitución española", |
|
theme="soft" |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
|
|
