app.py

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.chat_completion(inputs=inputs, task="text2text-generation")
    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 len(context) > 0:
        useful_context = context[0]  
        generation = generate_text(useful_context, query)
        return generation
    else:
        return "No se encontró información suficiente para responder."

    ranked = ranker.rerank(query=query, documents=raw_contexts, top_k=1)

    best_context = ranked[0]["content"]
    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()