Code walkthrough: Q&A on a custom PDF

Prerequisite

Make sure you've followed the steps in LLM application examples to clone our example repo, create your own OctoAI LLM endpoint, and set up your local environment.

Code walkthrough for answering questions about your custom PDF file

Below is an explanation of the code in https://github.com/octoml/octoml-llm-qa/blob/main/chat_main.py

First, we import the necessary libraries:

import logging  
import os  
import sys  
from pathlib import Path   

from dotenv import load_dotenv  
from OctoAiCloudLLM import OctoAiCloudLLM  
from langchain import HuggingFaceHub, OpenAI, PromptTemplate, LLMChain  
from langchain.embeddings import HuggingFaceEmbeddings  
from llama_index import ( LLMPredictor, ServiceContext, download_loader, GPTVectorStoreIndex, LangchainEmbedding)

We import libraries for environment variables, the OctoAI LLM endpoint, Langchain, embeddings, and the LLama index.

Next, we set the current directory and logging level:

# Get the current file's directory
current_dir = os.path.dirname(os.path.abspath(**file**))   

# Change the current working directory
os.chdir(current_dir)   

# Set logging level to CRITICAL
logging.basicConfig(level=logging.CRITICAL)

We need to load environment variables from a .env file to get credentials for the OctoAI model. We set the logging level to CRITICAL to reduce noise.

Then we define a function to initialize the files directory:

def init():  
    """Initialize the files directory."""  
    if not os.path.exists(FILES):  
        os.mkdir(FILES)

Next, we define a function to handle exiting the program:

def handle_exit():  
    """Handle exit gracefully."""  
    print("\nGoodbye!\n")  
    sys.exit(1)

The ask() function loads a PDF file, creates a query engine, and prompts the user to ask questions:

def ask(file):  
    """Load the file, create the query engine and interactively answer user questions about the document."""  
    print("Loading...")
    # Load the PDFReader
    PDFReader = download_loader("PDFReader")
    loader = PDFReader()
    documents = loader.load_data(file=Path(file))  

    # Initialize the OctoAiCloudLLM
    endpoint_url = os.getenv("ENDPOINT_URL")
    llm = OctoAiCloudLLM(endpoint_url=endpoint_url)
    llm_predictor = LLMPredictor(llm=llm)  

    # Create the LangchainEmbedding
    embeddings = LangchainEmbedding(
        HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2"))

    # Create the ServiceContext
    service_context = ServiceContext.from_defaults(
        llm_predictor=llm_predictor, chunk_size_limit=1024, embed_model=embeddings)  

    # Create the index from documents
    index = GPTVectorStoreIndex.from_documents(
        documents, service_context=service_context)

    # Create the query engine  
    query_engine = index.as_query_engine(
        verbose=True, llm_predictor=llm_predictor) 
    
    ...

We load the selected PDF, instantiate the OctoAI-hosted LLM and a predictor, create embeddings and a ServiceContext, build an index of the document, and construct a query engine to answer questions.

The select_file() function prompts the user to select a PDF file to process:

def select_file():  
    """Select a file for processing."""  
    ...  
    file_path = os.path.abspath(os.path.join(FILES, files[selection - 1]))  
    return file_path

Finally, we call the initialization function, prompt the user to select a file, and if a file is selected, start the interactive query session:

if __name__ == "__main__":
    # Initialize the file directory
    init()  

    # Prompt user to select a file 
    file = select_file()
    if file:
        # Start the interactive query session 
        ask(file)  
    else: 
        print("No files found")
        handle_exit()