QA AI Bot usando um LLM local e documentos

Neste tutorial vamos aprender a construir um Bot de Perguntas e Respostas (QA Bot) em Python usando um Large Language Model (LLM) local e documentos armazenados em uma pasta.

Pré-requisitos

Para começar, certifique-se de ter instalado as seguintes bibliotecas executando estes comandos no seu terminal:

pip install chromadb
pip install langchain
pip install transformers
pip install -U langchain-community

Carregar os documentos

Função para ler documentos

Vamos definir uma função que lê documentos de um diretório usando o DirectoryLoader.

from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import TextLoader

def load_docs(directory):
    loader = DirectoryLoader(directory, glob="**/*.txt", loader_cls=TextLoader)
    documents = loader.load()
    return documents

Depois, passamos o diretório para a função load_docs para carregar os documentos.

directory = 'Docs/'
documents = load_docs(directory)

Dividir os documentos em fragmentos

Para tornar o processamento mais eficiente, vamos dividir os documentos em fragmentos menores.

from langchain.text_splitter import RecursiveCharacterTextSplitter

def split_docs(documents, chunk_size=1000, chunk_overlap=20):
  text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
  docs = text_splitter.split_documents(documents)
  return docs

docs = split_docs(documents)

Inicializar os embeddings e o banco de dados vetorial

Inicializamos a classe SentenceTransformerEmbeddings e depois criamos um banco de dados vetorial usando o Chroma.

from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma

embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

Em seguida, inicializamos o Chroma com os documentos e os embeddings.

persist_directory = "zero"
vectordb = Chroma.from_documents(
    documents=docs, embedding=embeddings, persist_directory=persist_directory
)

Finalização do carregamento dos documentos

Carregamos os documentos do diretório persistente e obtemos um objeto retriever.

vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
vectordb.get()
retriever = vectordb.as_retriever(search_type="similarity", search_kwargs={"k": 2})

Criação do Bot de Perguntas e Respostas

Usaremos o pipeline do transformers para criar nosso QA Bot.

from transformers import pipeline

question_answerer = pipeline("question-answering", model='distilbert-base-cased-distilled-squad')

Formatar os documentos recuperados

Definimos uma função para formatar os documentos recuperados pelo retriever.

def format_docs(docs):
    if docs:
        return "\n\n".join(doc.page_content for doc in docs)

Interação com o bot

Imprimimos uma mensagem de boas-vindas e criamos um loop para interagir com o bot.

print("\033[92m" + "¡Bienvenido al Bot de QA AI ZEROGAP!" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\n")
while True:
    question = input("\033[93m" + "Tú: " + "\033[0m")
    print("\n")
    if question == "quit":
        break
    context = retriever.invoke(question)
    formated_context = format_docs(context)
    result = question_answerer(question=question, context=formated_context)

    print("\033[92m" + result['answer'] + "\033[0m")
    print("\n")

Código completo

A seguir, o código completo usado neste tutorial:

from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from transformers import pipeline

def load_docs(directory):
    loader = DirectoryLoader(directory, glob="**/*.txt", loader_cls=TextLoader)
    documents = loader.load()
    return documents

directory = 'Docs/'
documents = load_docs(directory)

def split_docs(documents, chunk_size=1000, chunk_overlap=20):
  text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
  docs = text_splitter.split_documents(documents)
  return docs

docs = split_docs(documents)

embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

persist_directory = "zero"
vectordb = Chroma.from_documents(
    documents=docs, embedding=embeddings, persist_directory=persist_directory
)

vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
vectordb.get()
retriever = vectordb.as_retriever(search_type="similarity", search_kwargs={"k": 2})

question_answerer = pipeline("question-answering", model='distilbert-base-cased-distilled-squad')

def format_docs(docs):
    if docs:
        return "\n\n".join(doc.page_content for doc in docs)

print("\033[92m" + "¡Bienvenido al Bot de QA AI ZEROGAP!" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\n")
while True:
    question = input("\033[93m" + "Tú: " + "\033[0m")
    print("\n")
    if question == "quit":
        break
    context = retriever.invoke(question)
    formated_context = format_docs(context)
    result = question_answerer(question=question, context=formated_context)

    print("\033[92m" + result['answer'] + "\033[0m")
    print("\n")

Com este código, você pode criar seu próprio Bot de Perguntas e Respostas usando documentos locais e um LLM. Experimente personalizá-lo e ampliá-lo de acordo com as necessidades do seu projeto!