QA AI Bot usando una LLM local y documentos

mayo 5, 2024

En este tutorial aprenderemos cómo construir un Bot de Preguntas y Respuestas (QA Bot) con python utilizando una Large Language Model (LLM) local y documentos almacenados en una carpeta.

Prerrequisitos

Para comenzar, asegúrate de haber instalado las siguientes librerías ejecutando estos comandos en tu terminal:

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

Cargar los documentos

Función para leer documentos

Vamos a definir una función que lee documentos desde un directorio utilizando 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

Después, pasaremos el directorio a la función load_docs para cargar los documentos.

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

Dividir los documentos en fragmentos

Para que el procesamiento sea más eficiente, dividiremos los documentos en fragmentos más pequeños.

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 las incrustaciones (embeddings) y la base de datos de vectores

Inicializamos la clase SentenceTransformerEmbeddings y luego creamos una base de datos de vectores utilizando Chroma.

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

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

Luego, inicializamos Chroma con los documentos y los embeddings.

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

Finalización de la carga de documentos

Cargamos documentos del directorio persistente y obtenemos un objeto retriever.

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

Creación del Bot de Preguntas y Respuestas

Usaremos pipeline de transformers para crear nuestro QA Bot.

from transformers import pipeline

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

Formatear los documentos recuperados

Definimos una función para formatear los documentos recuperados del retriever.

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

Interacción con el Bot

Imprimimos un mensaje de bienvenida y creamos un bucle para interactuar con el 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 continuación se muestra el código completo utilizado en este 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")

Con este código, puedes crear tu propio Bot de Preguntas y Respuestas utilizando documentos locales y una LLM. ¡Prueba a personalizarlo y expandirlo según las necesidades de tu proyecto!

Maximiliano Díaz Doglia
Últimas entradas de Maximiliano Díaz Doglia (ver todo)
Publicado en: AI