This repository contains two Python scripts that demonstrate how to build a document query and retrieval system using ChromaDB
as a vector store and HuggingFace
or Llama-CPP
models for embeddings and language modeling. The system leverages Box's text extraction capability to load documents, embeds them using a HuggingFace model, and queries them using a retrieval pipeline to provide relevant responses based on the input query.
It essentially follows the examples provide by Llama-Index in their documentation:
The repository includes two main scripts:
-
SimpleChroma.py: Simple Document Query System using ChromaDB
This script sets up a basic document query system usingChromaDB
for vector storage and a HuggingFace model for embeddings. It loads documents from Box, creates embeddings, stores them in a vector database, and uses an OpenAI model to query the database. -
NodesChroma.py: Advanced Document Query System with Llama-CPP
This script builds upon the first by adding a more advanced retrieval pipeline that leveragesLlama-CPP
for natural language processing and retrieval query synthesis. It includes features like text splitting, embedding generation, node construction, and an enhanced retrieval process.
To run the code, you can do the following:
python3.12 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
Before running the scripts, you need to configure the application settings:
-
Box Configuration
Ensure you have a Box developer account and set up an application for Client Credentials Grant (CCG). Common configurations include:- App + Enterprise Access
- Manage AI
- Generate user access tokens
- Remember to Authorize your app
-
API Keys and Environment Variables Copy the sample.env file to .env and update the values with your Box and OpenAI API keys.
- Steps:
- Load document data from Box using
BoxReaderTextExtraction
. - Convert document metadata to a compatible format.
- Set up the embedding model using
HuggingFaceEmbedding
. - Initialize ChromaDB as a vector store.
- Create an index from the documents using the embeddings.
- Query the index and print the results.
- Load document data from Box using
- Steps:
- Load document data from Box using
BoxReaderTextExtraction
. - Convert document metadata to a compatible format.
- Set up the embedding model using
HuggingFaceEmbedding
. - Initialize the LLM using
LlamaCPP
. - Use a text splitter to split documents into chunks.
- Construct nodes from the text chunks and generate embeddings.
- Load the nodes into a vector store.
- Build a retrieval pipeline and synthesize a response using
RetrieverQueryEngine
. - Query the index and print the results.
- Load document data from Box using
To run the scripts, execute them from the command line:
python SimpleChroma.py
Sample output:
================================================================================
Query: What to do in case of emergency?
--------------------------------------------------------------------------------
Query result:
If you become distressed on the surface during scuba diving, the appropriate action to take is to immediately drop your weight belt, inflate your buoyancy compensator (BC) for flotation, and signal for assistance.
--------------------------------------------------------------------------------
python NodesChroma.py
Sample Output:
================================================================================
Query: What equipment do I need to carry with me?
--------------------------------------------------------------------------------
Query result:
Based on the given context information, the answer to the query "What equipment do I need to carry with me?" is:
1. A snorkel
2. A submersible pressure gauge
3. A depth gauge
4. An alternate air source
5. A buoyancy compensator (BC) with an inflator hose and regulator.
This information can be found in the text under the section "Equipment Requirements".
--------------------------------------------------------------------------------
This project is licensed under the MIT License. See the LICENSE file for more details.
Contributions are welcome! Please open an issue or submit a pull request for any bugs or feature requests.