app.py

import gradio as gr
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFaceEndpoint
from langchain.memory import ConversationBufferMemory
from pathlib import Path
import chromadb
from unidecode import unidecode
import re

# Modelos LLM disponíveis
list_llm = [
    "mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.1",
    "google/gemma-7b-it", "google/gemma-2b-it", "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1",
    "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct",
    "tiiuae/falcon-7b-instruct", "google/flan-t5-xxl"
]
list_llm_simple = [os.path.basename(llm) for llm in list_llm]

# Função de carregamento e divisão de documentos
def load_and_split_documents(list_file_path, chunk_size, chunk_overlap):
    loaders = [PyPDFLoader(x) for x in list_file_path]
    pages = []
    for loader in loaders:
        pages.extend(loader.load())
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    return text_splitter.split_documents(pages)

# Função para criar banco de dados vetorial com ChromaDB
def create_vector_db(splits, collection_name):
    embedding = HuggingFaceEmbeddings()
    new_client = chromadb.PersistentClient(path="./chroma_db")
    return Chroma.from_documents(documents=splits, embedding=embedding, client=new_client, collection_name=collection_name)

# Função para inicializar a cadeia de QA
def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
    progress(0.1, desc="Inicializando tokenizer e Hub...")
    llm = HuggingFaceEndpoint(
        repo_id=llm_model, temperature=temperature, max_new_tokens=max_tokens, top_k=top_k, load_in_8bit=True
    )
    progress(0.5, desc="Definindo memória de buffer e cadeia de recuperação...")
    memory = ConversationBufferMemory(memory_key="chat_history", output_key='answer', return_messages=True)
    retriever = vector_db.as_retriever()
    qa_chain = ConversationalRetrievalChain.from_llm(llm, retriever=retriever, chain_type="stuff", memory=memory, return_source_documents=True)
    progress(0.9, desc="Concluído!")
    return qa_chain

# Função para gerar um nome de coleção válido
def create_collection_name(filepath):
    collection_name = Path(filepath).stem
    collection_name = unidecode(collection_name.replace(" ", "-"))
    return re.sub('[^A-Za-z0-9]+', '-', collection_name)[:50]

# Função para inicializar o banco de dados e o modelo LLM
def initialize_database_and_llm(list_file_obj, chunk_size, chunk_overlap, llm_option, llm_temperature, max_tokens, top_k, progress=gr.Progress()):
    list_file_path = [x.name for x in list_file_obj if x is not None]
    progress(0.1, desc="Criando nome da coleção...")
    collection_name = create_collection_name(list_file_path[0])
    progress(0.25, desc="Carregando e dividindo documentos...")
    doc_splits = load_and_split_documents(list_file_path, chunk_size, chunk_overlap)
    progress(0.5, desc="Gerando banco de dados vetorial...")
    vector_db = create_vector_db(doc_splits, collection_name)
    progress(0.75, desc="Inicializando modelo LLM...")
    llm_name = list_llm[llm_option]
    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
    progress(0.9, desc="Concluído!")
    return vector_db, collection_name, qa_chain

# Função de interação com o chatbot
def conversation(qa_chain, message, history):
    formatted_chat_history = [f"Usuário: {user_message}\nAssistente: {bot_message}" for user_message, bot_message in history]
    response = qa_chain({"question": message, "chat_history": formatted_chat_history})
    response_answer = response["answer"].split("Resposta útil:")[-1]
    response_sources = [doc.page_content.strip() for doc in response["source_documents"]]
    response_pages = [doc.metadata["page"] + 1 for doc in response["source_documents"]]
    new_history = history + [(message, response_answer)]
    return qa_chain, gr.update(value=""), new_history, *response_sources, *response_pages

# Função de carregamento de arquivos
def upload_file(file_obj):
    return [file_obj.name for file_obj in file_obj if file_obj is not None]

# Interface Gradio
def demo():
    with gr.Blocks(theme="base") as demo:
        vector_db, qa_chain, collection_name = gr.State(), gr.State(), gr.State()
        gr.Markdown("<center><h2>Chatbot baseado em PDF</center></h2><h3>Faça qualquer pergunta sobre seus documentos PDF</h3>")

        with gr.Tab("Etapa 1 - Carregar PDF"):
            document = gr.Files(height=100, file_count="multiple", file_types=["pdf"])

        with gr.Tab("Etapa 2 - Processar documento"):
            db_btn = gr.Button("Gerar banco de dados vetorial")
            slider_chunk_size = gr.Slider(minimum=100, maximum=1000, value=600, step=20, label="Tamanho do bloco")
            slider_chunk_overlap = gr.Slider(minimum=10, maximum=200, value=40, step=10, label="Sobreposição do bloco")
            db_progress = gr.Textbox(label="Inicialização do banco de dados vetorial")

        with gr.Tab("Etapa 3 - Inicializar cadeia de QA"):
            llm_btn = gr.Radio(list_llm_simple, label="Modelos LLM")
            slider_temperature = gr.Slider(minimum=0.01, maximum=1.0, value=0.7, step=0.1, label="Temperatura")
            slider_maxtokens = gr.Slider(minimum=224, maximum=4096, value=1024, step=32, label="Máximo de Tokens")
            slider_topk = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Amostras top-k")
            llm_progress = gr.Textbox(value="Nenhum", label="Inicialização da cadeia QA")
            qachain_btn = gr.Button("Inicializar cadeia de Pergunta e Resposta")

        with gr.Tab("Etapa 4 - Chatbot"):
            chatbot = gr.Chatbot(height=300)
            doc_source1, doc_source2, doc_source3 = gr.Textbox(label="Referência 1"), gr.Textbox(label="Referência 2"), gr.Textbox(label="Referência 3")
            source1_page, source2_page, source3_page = gr.Number(label="Página 1"), gr.Number(label="Página 2"), gr.Number(label="Página 3")
            
            # Campo de texto para enviar mensagens
            user_input = gr.Textbox(label="Sua mensagem")

            # Implementação de lógica de interação de conversa
            def send_message(message, history, qa_chain):
                formatted_chat_history = [f"Usuário: {user_message}\nAssistente: {bot_message}" for user_message, bot_message in history]
                response = qa_chain({"question": message, "chat_history": formatted_chat_history})
                response_answer = response["answer"].split("Resposta útil:")[-1]
                response_sources = [doc.page_content.strip() for doc in response["source_documents"]]
                response_pages = [doc.metadata["page"] + 1 for doc in response["source_documents"]]
                new_history = history + [(message, response_answer)]
                return qa_chain, gr.update(value=""), new_history, *response_sources, *response_pages

            user_input.submit(send_message, inputs=[user_input, chatbot.history, qa_chain], outputs=[qa_chain, gr.update(value=""), chatbot.history, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page])
        
    demo.launch()