Skip to content

RAG -Search

Install packages

1
2
3
4
5
6
7
!uv pip install -q \
    pandas==2.3.2 \
    pandas-stubs==2.3.2.250827 \
    numpy==2.3.2 \
    matplotlib==3.10.6 \
    scikit-learn==1.7.1 \
    requests==2.32.5

Import packages

1
2
3
4
5
6
import numpy as np
import pandas as pd
import requests
from sklearn.decomposition import NMF, TruncatedSVD
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

Get documents

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
docs_url = "https://github.com/alexeygrigorev/llm-rag-workshop/raw/main/notebooks/documents.json"
docs_response = requests.get(docs_url)
documents_raw = docs_response.json()

documents = []

for course in documents_raw:
    course_name = course["course"]
    for doc in course["documents"]:
        doc["course"] = course_name
        documents.append(doc)

documents[0]

{'text': "The purpose of this document is to capture frequently asked technical questions\nThe exact day and hour of the course will be 15th Jan 2024 at 17h00. The course will start with the first  “Office Hours'' live.1\nSubscribe to course public Google Calendar (it works from Desktop only).\nRegister before the course starts using this link.\nJoin the course Telegram channel with announcements.\nDon’t forget to register in DataTalks.Club's Slack and join the channel.",
 'section': 'General course-related questions',
 'question': 'Course - When will the course start?',
 'course': 'data-engineering-zoomcamp'}

Create a DataFrame

1
2
df = pd.DataFrame(documents, columns=["course", "section", "question", "text"])
df.head()
course section question text
0 data-engineering-zoomcamp General course-related questions Course - When will the course start? The purpose of this document is to capture fre...
1 data-engineering-zoomcamp General course-related questions Course - What are the prerequisites for this c... GitHub - DataTalksClub data-engineering-zoomca...
2 data-engineering-zoomcamp General course-related questions Course - Can I still join the course after the... Yes, even if you don't register, you're still ...
3 data-engineering-zoomcamp General course-related questions Course - I have registered for the Data Engine... You don't need it. You're accepted. You can al...
4 data-engineering-zoomcamp General course-related questions Course - What can I do before the course starts? You can start by installing and setting up all...

Vector spaces

  • Turn docs into vectors
  • Term-document-matrix:
  • rows: documents
  • columns: words/tokens
  • Bag of words: counts occurrences of each word in each document
Document word1 word2 word3
Doc1 1 0 2
Doc2 0 3 1 
1
2
3
4
cv = CountVectorizer(min_df=5, stop_words="english")  # Min Document Frequency
X = cv.fit_transform(df.text)
column_names = cv.get_feature_names_out()
column_names

array(['01', '02', '03', ..., 'youtube', 'zip', 'zoomcamp'],
      shape=(1333,), dtype=object)

Number of distinct tokens

1
cv.get_feature_names_out().shape

(1333,)

Tokens by occurrences

1
2
df_docs = pd.DataFrame(X.toarray(), columns=column_names)
df_docs.head()
01 02 03 04 05 06 09 10 100 11 ... y_val yaml year yellow yellow_tripdata_2021 yes yml youtube zip zoomcamp
0 0 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 0 0
1 0 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 0 1
2 0 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 1 0 0 0 0
3 0 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 0 0
4 0 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 0 0

5 rows × 1333 columns

Tokens by importance

1
2
3
4
5
6
tv = TfidfVectorizer(min_df=5, stop_words="english")  # Min Document Frequency
X = tv.fit_transform(df.text)
column_names = tv.get_feature_names_out()
column_names
df_docs = pd.DataFrame(X.toarray(), columns=column_names)
df_docs.head()
01 02 03 04 05 06 09 10 100 11 ... y_val yaml year yellow yellow_tripdata_2021 yes yml youtube zip zoomcamp
0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.000000 0.0 0.0 0.0 0.000000
1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.000000 0.0 0.0 0.0 0.428961
2 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.279891 0.0 0.0 0.0 0.000000
3 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.000000 0.0 0.0 0.0 0.000000
4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.000000 0.0 0.0 0.0 0.000000

5 rows × 1333 columns

Text Search

1
2
3
4
query = "I just discovered the course, is it too late to join?"

q = cv.transform([query])
q.toarray()

array([[0, 0, 0, ..., 0, 0, 0]], shape=(1, 1333))

Relevant documents

1
2
score = cosine_similarity(X, q).flatten()
np.argsort(score)[-5:]

array([ 27, 449, 440,  22,   0])

Document example

1
df.iloc[0].text

"The purpose of this document is to capture frequently asked technical questions\nThe exact day and hour of the course will be 15th Jan 2024 at 17h00. The course will start with the first  “Office Hours'' live.1\nSubscribe to course public Google Calendar (it works from Desktop only).\nRegister before the course starts using this link.\nJoin the course Telegram channel with announcements.\nDon’t forget to register in DataTalks.Club's Slack and join the channel."

Creating matrices

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
fields = ["section", "question", "text"]

matrices = {}
vectorizers = {}

for f in fields:
    cv = TfidfVectorizer(stop_words="english", min_df=5)
    X = cv.fit_transform(df[f])
    matrices[f] = X
    vectorizers[f] = cv

Scoring

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
n = len(df)
score = np.zeros(n)
boosts = {"question": 3}

for f in fields:
    q = vectorizers[f].transform([query])
    X = matrices[f]

    f_score = cosine_similarity(X, q).flatten()
    boost = boosts.get(f, 1.0)
    score = score + boost * f_score

Filtering

1
2
3
4
5
filters = {"course": "data-engineering-zoomcamp"}

for field, value in filters.items():
    mask = (df[field] == value).values
    score = score * mask

Getting most relevant documents

1
2
idx = np.argsort(score)[-5:][::-1]
df.iloc[idx]
course section question text
7 data-engineering-zoomcamp General course-related questions Course - Can I follow the course after it fini... Yes, we will keep all the materials after the ...
0 data-engineering-zoomcamp General course-related questions Course - When will the course start? The purpose of this document is to capture fre...
1 data-engineering-zoomcamp General course-related questions Course - What are the prerequisites for this c... GitHub - DataTalksClub data-engineering-zoomca...
34 data-engineering-zoomcamp General course-related questions How can we contribute to the course? Star the repo! Share it with friends if you fi...
4 data-engineering-zoomcamp General course-related questions Course - What can I do before the course starts? You can start by installing and setting up all...

Putting all together

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class TextSearch:
    def __init__(self, text_fields):
        self.text_fields = text_fields
        self.matrices = {}
        self.vectorizers = {}

    def fit(self, records, vectorizer_params={}):
        self.df = pd.DataFrame(records)

        for f in self.text_fields:
            cv = TfidfVectorizer(**vectorizer_params)
            X = cv.fit_transform(self.df[f])
            self.matrices[f] = X
            self.vectorizers[f] = cv

    def search(self, query, n_results=10, boost={}, filters={}):
        score = np.zeros(len(self.df))

        for f in self.text_fields:
            b = boost.get(f, 1.0)
            q = self.vectorizers[f].transform([query])
            s = cosine_similarity(self.matrices[f], q).flatten()
            score = score + b * s

        for field, value in filters.items():
            mask = (self.df[field] == value).values
            score = score * mask

        idx = np.argsort(-score)[:n_results]
        results = self.df.iloc[idx]
        return results.to_dict(orient="records")

Searching

1
2
3
4
5
6
7
8
9
index = TextSearch(text_fields=["section", "question", "text"])
index.fit(documents)

index.search(
    query="I just signed up. Is it too late to join the course?",
    n_results=5,
    boost={"question": 3.0},
    filters={"course": "data-engineering-zoomcamp"},
)

[{'text': "Yes, even if you don't register, you're still eligible to submit the homeworks.\nBe aware, however, that there will be deadlines for turning in the final projects. So don't leave everything for the last minute.",
  'section': 'General course-related questions',
  'question': 'Course - Can I still join the course after the start date?',
  'course': 'data-engineering-zoomcamp'},
 {'text': "The purpose of this document is to capture frequently asked technical questions\nThe exact day and hour of the course will be 15th Jan 2024 at 17h00. The course will start with the first  “Office Hours'' live.1\nSubscribe to course public Google Calendar (it works from Desktop only).\nRegister before the course starts using this link.\nJoin the course Telegram channel with announcements.\nDon’t forget to register in DataTalks.Club's Slack and join the channel.",
  'section': 'General course-related questions',
  'question': 'Course - When will the course start?',
  'course': 'data-engineering-zoomcamp'},
 {'text': 'Yes, we will keep all the materials after the course finishes, so you can follow the course at your own pace after it finishes.\nYou can also continue looking at the homeworks and continue preparing for the next cohort. I guess you can also start working on your final capstone project.',
  'section': 'General course-related questions',
  'question': 'Course - Can I follow the course after it finishes?',
  'course': 'data-engineering-zoomcamp'},
 {'text': 'No, late submissions are not allowed. But if the form is still not closed and it’s after the due date, you can still submit the homework. confirm your submission by the date-timestamp on the Course page.y\nOlder news:[source1] [source2]',
  'section': 'General course-related questions',
  'question': 'Homework - Are late submissions of homework allowed?',
  'course': 'data-engineering-zoomcamp'},
 {'text': 'You can start by installing and setting up all the dependencies and requirements:\nGoogle cloud account\nGoogle Cloud SDK\nPython 3 (installed with Anaconda)\nTerraform\nGit\nLook over the prerequisites and syllabus to see if you are comfortable with these subjects.',
  'section': 'General course-related questions',
  'question': 'Course - What can I do before the course starts?',
  'course': 'data-engineering-zoomcamp'}]

1
2
3
4
5
6
7
X = matrices["text"]
cv = vectorizers["text"]

svd = TruncatedSVD(n_components=16)
X_emb = svd.fit_transform(X)

X_emb.shape

(948, 16)

Dense representation of documents (embeddings)

1
X_emb[0]

array([ 0.09652939, -0.08220536, -0.10269583, -0.07888492,  0.0688845 ,
       -0.06010652,  0.01896283, -0.15131328, -0.22164436, -0.27420628,
        0.12057112,  0.06016185, -0.10346335,  0.09990854,  0.03095138,
       -0.01011112])

1
2
3
4
5
query = "I just signed up. Is it too late to join the course?"

Q = cv.transform([query])
Q_emb = svd.transform(Q)
Q_emb[0]

array([ 0.05790068, -0.03850882, -0.05679924, -0.0280689 ,  0.04099817,
       -0.06254927,  0.00922403, -0.09846986, -0.15605146, -0.18274481,
        0.09566324,  0.06353763, -0.07463788,  0.05534739,  0.03201764,
       -0.01498711])

1
np.dot(X_emb[0], Q_emb[0])

np.float64(0.15290943169419263)

1
2
3
score = cosine_similarity(X_emb, Q_emb).flatten()
idx = np.argsort(-score)[:10]
df.loc[idx]
course section question text
764 machine-learning-zoomcamp Projects (Midterm and Capstone) What If I submitted only two projects and fail... If you have submitted two projects (and peer-r...
0 data-engineering-zoomcamp General course-related questions Course - When will the course start? The purpose of this document is to capture fre...
449 machine-learning-zoomcamp General course-related questions The course has already started. Can I still jo... Yes, you can. You won’t be able to submit some...
451 machine-learning-zoomcamp General course-related questions Can I submit the homework after the due date? No, it’s not possible. The form is closed afte...
814 mlops-zoomcamp +-General course questions What if my answer is not exactly the same as t... Please choose the closest one to your answer. ...
2 data-engineering-zoomcamp General course-related questions Course - Can I still join the course after the... Yes, even if you don't register, you're still ...
436 machine-learning-zoomcamp General course-related questions Is it going to be live? When? The course videos are pre-recorded, you can st...
13 data-engineering-zoomcamp General course-related questions Office Hours - I can’t attend the “Office hour... Yes! Every “Office Hours” will be recorded and...
7 data-engineering-zoomcamp General course-related questions Course - Can I follow the course after it fini... Yes, we will keep all the materials after the ...
15 data-engineering-zoomcamp General course-related questions Homework - Are late submissions of homework al... No, late submissions are not allowed. But if t... 
1
2
3
nmf = NMF(n_components=16)
X_emb = nmf.fit_transform(X)
X_emb[0]

array([0.        , 0.        , 0.        , 0.        , 0.        ,
       0.        , 0.        , 0.        , 0.27529579, 0.        ,
       0.        , 0.        , 0.        , 0.        , 0.        ,
       0.        ])

1
2
3
Q = cv.transform([query])
Q_emb = nmf.transform(Q)
Q_emb[0]

array([0.        , 0.00224623, 0.        , 0.        , 0.        ,
       0.        , 0.        , 0.        , 0.1818919 , 0.        ,
       0.        , 0.        , 0.        , 0.        , 0.0017897 ,
       0.        ])

1
2
3
score = cosine_similarity(X_emb, Q_emb).flatten()
idx = np.argsort(-score)[:10]
df.loc[idx]
course section question text
2 data-engineering-zoomcamp General course-related questions Course - Can I still join the course after the... Yes, even if you don't register, you're still ...
0 data-engineering-zoomcamp General course-related questions Course - When will the course start? The purpose of this document is to capture fre...
436 machine-learning-zoomcamp General course-related questions Is it going to be live? When? The course videos are pre-recorded, you can st...
449 machine-learning-zoomcamp General course-related questions The course has already started. Can I still jo... Yes, you can. You won’t be able to submit some...
814 mlops-zoomcamp +-General course questions What if my answer is not exactly the same as t... Please choose the closest one to your answer. ...
764 machine-learning-zoomcamp Projects (Midterm and Capstone) What If I submitted only two projects and fail... If you have submitted two projects (and peer-r...
451 machine-learning-zoomcamp General course-related questions Can I submit the homework after the due date? No, it’s not possible. The form is closed afte...
437 machine-learning-zoomcamp General course-related questions What if I miss a session? Everything is recorded, so you won’t miss anyt...
11 data-engineering-zoomcamp General course-related questions Certificate - Can I follow the course in a sel... No, you can only get a certificate if you fini...
7 data-engineering-zoomcamp General course-related questions Course - Can I follow the course after it fini... Yes, we will keep all the materials after the ...