QA AI Bot using a local LLM and documents

In this tutorial we’ll learn how to build a Question-and-Answer Bot (QA Bot) in Python using a local Large Language Model (LLM) and documents stored in a folder.

Prerequisites

To get started, make sure you’ve installed the following libraries by running these commands in your terminal:

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

Load the documents

Function to read documents

We’ll define a function that reads documents from a directory using 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

Then we pass the directory to the load_docs function to load the documents.

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

Split the documents into chunks

To make processing more efficient, we’ll split the documents into smaller chunks.

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)

Initialise the embeddings and the vector database

We initialise the SentenceTransformerEmbeddings class and then create a vector database using Chroma.

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

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

Next, we initialise Chroma with the documents and the embeddings.

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

Finishing the document load

We load documents from the persistent directory and get a retriever object.

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

Building the Question-and-Answer Bot

We’ll use the pipeline from transformers to create our QA Bot.

from transformers import pipeline

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

Format the retrieved documents

We define a function to format the documents returned by the retriever.

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

Interacting with the Bot

We print a welcome message and create a loop to interact with the 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")

Complete code

Below is the complete code used in this 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")

With this code you can build your own Question-and-Answer Bot using local documents and an LLM. Try customising and extending it to fit your project’s needs!