Document Indexing and RAG

import os
import getpass
if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')

%pip install -q pixeltable sentence-transformers tiktoken openai openpyxl

Note: you may need to restart the kernel to use updated packages.

import numpy as np
import pixeltable as pxt

# Ensure a clean slate for the demo
pxt.drop_dir('rag_demo', force=True)
pxt.create_dir('rag_demo')

Connected to Pixeltable database at: postgresql+psycopg://postgres:@/pixeltable?host=/home/marcel/.pixeltable/pgdata
Created directory 'rag_demo'.

<pixeltable.catalog.dir.Dir at 0x7eb34c1eceb0>

base = 'https://github.com/pixeltable/pixeltable/raw/main/docs/resources/rag-demo/'
qa_url = base + 'Q-A-Rag.xlsx'
queries_t = pxt.create_table('rag_demo.queries', source=qa_url)

Created table 'queries'.
Inserting rows into `queries`: 8 rows [00:00, 1853.94 rows/s]
Inserted 8 rows with 0 errors.

queries_t.head()

documents_t = pxt.create_table(
    'rag_demo.documents',
    {'document': pxt.Document}
)

documents_t

Created table 'documents'.

document_urls = [
    base + 'Argus-Market-Digest-June-2024.pdf',
    base + 'Argus-Market-Watch-June-2024.pdf',
    base + 'Company-Research-Alphabet.pdf',
    base + 'Jefferson-Amazon.pdf',
    base + 'Mclean-Equity-Alphabet.pdf',
    base + 'Zacks-Nvidia-Report.pdf',
]

documents_t.insert({'document': url} for url in document_urls[:3])
documents_t.show()

Inserting rows into `documents`: 3 rows [00:00, 1925.76 rows/s]
Inserted 3 rows with 0 errors.

from pixeltable.iterators import DocumentSplitter

chunks_t = pxt.create_view(
    'rag_demo.chunks',
    documents_t,
    iterator=DocumentSplitter.create(
        document=documents_t.document,
        separators='token_limit',
        limit=300
    )
)

Inserting rows into `chunks`: 41 rows [00:00, 21767.91 rows/s]

chunks_t

chunks_t.where(chunks_t.pos < 2).show()

from pixeltable.functions.huggingface import sentence_transformer

chunks_t.add_embedding_index(
    'text',
    embedding=sentence_transformer.using(model_id='intfloat/e5-large-v2')
)

query_text = "What is the expected EPS for Nvidia in Q1 2026?"
sim = chunks_t.text.similarity(query_text)
nvidia_eps_query = (
    chunks_t
    .order_by(sim, asc=False)
    .select(similarity=sim, text=chunks_t.text)
    .limit(5)
)
nvidia_eps_query.collect()

# A @query is essentially a reusable, parameterized query that is attached to a table (or view),
# which is a modular way of getting data from that table.

@pxt.query
def top_k(query_text: str):
    sim = chunks_t.text.similarity(query_text)
    return (
        chunks_t.order_by(sim, asc=False)
            .select(chunks_t.text, sim=sim)
            .limit(5)
    )

# Now add a computed column to `queries_t`, calling the query
# `top_k` that we just defined.
queries_t.add_computed_column(
    question_context=top_k(queries_t.Question)
)

Added 8 column values with 0 errors.

8 rows updated, 8 values computed.

queries_t

queries_t.select(queries_t.question_context).head(1)

# Define a UDF to create an LLM prompt given a top-k list of
# context chunks and a question.
@pxt.udf
def create_prompt(top_k_list: list[dict], question: str) -> str:
    concat_top_k = '\n\n'.join(
        elt['text'] for elt in reversed(top_k_list)
    )
    return f'''
    PASSAGES:

    {concat_top_k}

    QUESTION:

    {question}'''

queries_t.add_computed_column(
    prompt=create_prompt(queries_t.question_context, queries_t.Question)
)

Added 8 column values with 0 errors.

8 rows updated, 16 values computed.

queries_t

queries_t.select(queries_t.prompt).head(1)

from pixeltable.functions import openai

# Assemble the prompt and instructions into OpenAI's message format
messages = [
    {
        'role': 'system',
        'content': 'Please read the following passages and answer the question based on their contents.'
    },
    {
        'role': 'user',
        'content': queries_t.prompt
    }
]

# Add a computed column that calls OpenAI
queries_t.add_computed_column(
    response=openai.chat_completions(model='gpt-4o-mini', messages=messages)
)

Added 8 column values with 0 errors.

8 rows updated, 8 values computed.

queries_t.add_computed_column(
    answer=queries_t.response.choices[0].message.content
)

Added 8 column values with 0 errors.

8 rows updated, 8 values computed.

queries_t

queries_t.select(queries_t.Question, queries_t.correct_answer, queries_t.answer).show()

documents_t.insert({'document': p} for p in document_urls[3:])

Inserting rows into `documents`: 3 rows [00:00, 1949.63 rows/s]
Inserting rows into `chunks`: 68 rows [00:00, 601.57 rows/s]
Inserted 71 rows with 0 errors.

71 rows inserted, 6 values computed.

documents_t.show()

nvidia_eps_query.collect()

queries_t.recompute_columns('response')

Inserting rows into `queries`: 8 rows [00:00, 2128.01 rows/s]

8 rows updated, 16 values computed.

queries_t.select(
    queries_t.Question,
    queries_t.correct_answer,
    queries_t.answer
).show()