File size: 1,999 Bytes
1b3bcd8
9a0264a
 
 
 
 
 
 
1b3bcd8
9a0264a
 
1b3bcd8
9a0264a
 
1b3bcd8
9a0264a
 
1b3bcd8
9a0264a
1b3bcd8
9a0264a
1b3bcd8
9a0264a
1b3bcd8
9a0264a
1b3bcd8
9a0264a
 
a06c139
9a0264a
a06c139
9a0264a
 
 
cf4a4dc
9a0264a
cf4a4dc
9a0264a
 
1b3bcd8
a06c139
9a0264a
 
cf4a4dc
 
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
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

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)

embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

vectordb = Chroma.from_documents(docs_split, embedding_function)

client = InferenceClient("google/flan-t5-base", token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))

ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type='colbert')

def generate_text(context, query):
    inputs = f"Context: {context} Question: {query}"
    response = client.text_generation(prompt=inputs)
    return response
    
def test_rag_reranking(query, ranker):
    docs = vectordb.similarity_search_with_score(query)
    context = []

    for doc, score in docs:
        if score < 7:
            doc_details = doc.to_json()['kwargs']
            context.append(doc_details['page_content'])

    if not context:
        return "No se encontr贸 informaci贸n suficiente para responder."

    reranked = ranker.rank(query=query, documents=context, top_k=1)
    best_context = reranked[0]["text"]  

    return generate_text(best_context, query)

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()