File Search

This tutorial demonstrates how to create a vector index for PDF files, and search over the files using an English phrase.
How to run this tutorial

We recommend running this tutorial, as well as the other various tutorials in the Tutorials section, inside TileDB Cloud. This will allow you to quickly experiment avoiding all the installation, deployment, and configuration hassles. Sign up for the free tier, spin up a TileDB Cloud notebook with a Python kernel, and follow the tutorial instructions. If you wish to learn how to run tutorials locally on your machine, read the Tutorials: Running Locally tutorial.

In this tutorial, you will learn how to load large collections of PDF files into a TileDB-Vector-Search index, and query them using an English phrase.

Setup

To be able to run this tutorial, you will need an OpenAI API key. In addition, if you wish to use your local machine instead of a TileDB Cloud notebook, you will need to install the following:

pip install langchain==0.0.331 langchain_community openai==0.28.1 tiktoken pymupdf pillow

Start by importing the necessary libraries, setting the URIs you will use throughout the tutorial, and clean up any previously generated data.

import os
import shutil

os.environ["TOKENIZERS_PARALLELISM"] = "true"
import warnings

warnings.filterwarnings("ignore")
import numpy as np
import fitz
from PIL import Image
from io import BytesIO

import tiledb
import tiledb.vector_search as vs
from langchain.document_loaders.generic import GenericLoader
from langchain_community.document_loaders.parsers.pdf import PyMuPDFParser
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings

# URIs you will use in this tutorial
input_files_uri = "random_invoices"
index_uri = "file_index"
metadata_array_uri = "chunk_metadata"

# Clean up past data
if os.path.exists(input_files_uri):
    shutil.rmtree(input_files_uri)
if os.path.exists(index_uri):
    shutil.rmtree(index_uri)
if os.path.exists(metadata_array_uri):
    shutil.rmtree(metadata_array_uri)

Download some synthetically generated invoice PDFs.

os.mkdir(input_files_uri)
!aws s3 cp --no-sign-request --recursive s3://tiledb-inc-demo-data/examples/notebooks/genai/random_invoices {input_files_uri}

Ingestion

You will ingest the invoice PDFs. The ingestion performs:

  • File parsing, text extraction, and text splitting into chunks.
  • Text embedding generation using open source embedding models or OpenAI API calls.
  • Vector indexing of embeddings.

Extract text from the PDF documents and split it into text chunks using LangChain utilities.

# Parse documents and split them into text chunks
loader = GenericLoader.from_filesystem(
    input_files_uri, glob="**/*", suffixes=[".pdf"], parser=PyMuPDFParser()
)
documents = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
documents = loader.load()
print(f"Number of raw documents loaded: {len(documents)}")
documents = splitter.split_documents(documents)
texts = [d.page_content for d in documents]
print(f"Number of document chunks: {len(texts)}")
Number of raw documents loaded: 971
Number of document chunks: 971

Create a TileDB array to store the chunks.

# Create metadata array
external_ids_dim = tiledb.Dim(
    name="external_id",
    domain=(0, np.iinfo(np.dtype("uint64")).max - 10000),
    tile=10000,
    dtype=np.dtype(np.uint64),
)
external_ids_dom = tiledb.Domain(external_ids_dim)
text_attr = tiledb.Attr(name="text", dtype=str)
file_path_attr = tiledb.Attr(name="file_path", dtype=str)
page_attr = tiledb.Attr(name="page", dtype=np.int32)
attrs = [text_attr, file_path_attr, page_attr]
schema = tiledb.ArraySchema(
    domain=external_ids_dom,
    sparse=True,
    attrs=attrs,
)
tiledb.Array.create(metadata_array_uri, schema)

Next, store the text of the chunks along with their metadata in the TileDB array.

# Add document chunk metadata to metadata array
size = len(documents)
text_metadata = np.empty(size, dtype="O")
file_paths = np.empty(size, dtype="O")
pages = np.zeros(size, dtype=np.int32)
external_ids = np.zeros(size, dtype=np.uint64)

for i in range(size):
    pages[i] = int(documents[i].metadata["page"])
    file_paths[i] = documents[i].metadata["file_path"]
    text_metadata[i] = texts[i]
    external_ids[i] = i

with tiledb.open(metadata_array_uri, "w") as metadata_array:
    metadata_array[external_ids] = {
        "text": text_metadata,
        "file_path": file_paths,
        "page": pages,
    }

You can now generate text embeddings and index them using an IVF_FLAT index.

# NOTE: You need to set the OPENAI_API_KEY variable for this to work.

# Generate embeddings for each document chunk
embedding = OpenAIEmbeddings()
text_embeddings = embedding.embed_documents(texts)

# Index document chunks using a TileDB IVF_FLAT index
vs.ingest(
    index_type="IVF_FLAT",
    index_uri=index_uri,
    input_vectors=np.array(text_embeddings).astype(np.float32),
)
print(
    f"Number of vector embeddings stored in TileDB-Vector-Search: {len(text_embeddings)}"
)
Number of vector embeddings stored in TileDB-Vector-Search: 971

Clean up

Delete the data you created in this tutorial.

# Clean up
if os.path.exists(input_files_uri):
    shutil.rmtree(input_files_uri)
if os.path.exists(index_uri):
    shutil.rmtree(index_uri)
if os.path.exists(metadata_array_uri):
    shutil.rmtree(metadata_array_uri)