File size: 3,100 Bytes
1b3bcd8
9a0264a
 
 
 
 
 
 
1b3bcd8
9a0264a
 
1b3bcd8
af2fcf7
 
1b3bcd8
af2fcf7
 
 
1b3bcd8
af2fcf7
 
1b3bcd8
af2fcf7
 
 
 
1b3bcd8
551b707
1b3bcd8
9a0264a
1b3bcd8
9a0264a
5c340bd
 
 
 
 
 
 
ccd7fd0
 
a06c139
9a0264a
b9eeb67
 
9a0264a
b9eeb67
cf4a4dc
b9eeb67
 
 
cf4a4dc
9a0264a
b9eeb67
 
 
1b3bcd8
a06c139
b9eeb67
9a0264a
 
b9eeb67
 
946d2c8
61bee11
5dd204e
 
b9eeb67
 
7c37226
d817f91
b9eeb67
 
 
 
 
9a0264a
 
 
 
1b3bcd8
 
9a0264a
 
 
1b3bcd8
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
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]
    #print(help(Reranker.models.ColBERTRanker.rank))
    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()