##########################################################################################################
'''
The following script translates the list of strings "text_chunks" into vector embeddings and
saves the DataFrame including "text_chunks" and "embeddings" in the csv file "embeddings_df.csv"
'''
##########################################################################################################

import os
import requests
import pandas as pd
import numpy as np

# hugging face token
os.environ['hf_token'] = os.environ.get('HF_TOKEN')
# os.environ['hf_token'] = 'testtoken123'

# example text snippets we want to translate into vector embeddings
text_chunks = [
    "The sky is blue.",
    "The grass is green.",
    "The sun is shining.",
    "I love chocolate.",
    "Pizza is delicious.",
    "Coding is fun.",
    "Roses are red.",
    "Violets are blue.",
    "Water is essential for life.",
    "The moon orbits the Earth.",
]

def _get_embeddings(text_chunk):
    '''
    Use embedding model from hugging face to calculate embeddings for the text snippets provided
   
    Parameters:
        - text_chunk (string): the sentence or text snippet you want to translate into embeddings

    Returns:
        - embedding(list): list with all embedding dimensions
    '''
    # define the embedding model you want to use
    model_id = "sentence-transformers/all-MiniLM-L6-v2"

    # you can find the token to the hugging face api in your settings page https://huggingface.co/settings/tokens
    hf_token = os.environ.get('hf_token')

    # API endpoint for embedding model
    api_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model_id}"
    headers = {"Authorization": f"Bearer {hf_token}"}

    # call API
    response = requests.post(api_url, headers=headers, json={"inputs": text_chunk, "options":{"wait_for_model":True}})
   
    # load response from embedding model into json format
    embedding = response.json()

    return embedding

def from_text_to_embeddings(text_chunks):
    '''
    Translate sentences into vector embeddings

    Attributes:
        - text_chunks (list): list of example strings

    Returns:
        - embeddings_df (DataFrame): data frame with the columns "text_chunk" and "embeddings"
    '''
    # create new data frame using text chunks list
    embeddings_df = pd.DataFrame(text_chunks).rename(columns={0:"text_chunk"})

    # use the _get_embeddings function to retrieve the embeddings for each of the sentences
    embeddings_df["embeddings"] = embeddings_df["text_chunk"].apply(_get_embeddings)

    # split the embeddings column into individuell columns for each vector dimension
    embeddings_df = embeddings_df['embeddings'].apply(pd.Series)
    embeddings_df["text_chunk"] = text_chunks

    return embeddings_df

# get embeddings for each of the text chunks
embeddings_df = from_text_to_embeddings(text_chunks)

# save data frame with text chunks and embeddings to csv
embeddings_df.to_csv('../02_Data/embeddings_df.csv', index=False)
01_Data_Pipelines/01_text_to_embeddings.py
polzerdo55862/vector-store-tutorial
##########################################################################################################
# Create PCA plot using embeddings df
##########################################################################################################
import pandas as pd
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

def create_pca_plot(embeddings_df):
    '''
    The function performs a principal component analysis to reduce the dimensions to 2 so we can print them in a plot

    Parameters
        - embeddings_df (DataFrame): data frame with the columns "text_chunk" and "embeddings"

    Returns
        - df_reduced (DataFrame): data frame with the 2 most relevant Principal Components
    '''
    # Perform PCA with 2 components
    pca = PCA(n_components=2)

    # apply principal component analysis to the embeddings table
    df_reduced = pca.fit_transform(embeddings_df[embeddings_df.columns[:-2]])

    # Create a new DataFrame with reduced dimensions
    df_reduced = pd.DataFrame(df_reduced, columns=['PC1', 'PC2'])

    ############################################################################################
    # Create a scatter plot
    ############################################################################################
    def create_scatter_plot(df_reduced):
        plt.scatter(df_reduced['PC1'], df_reduced['PC2'], label=df_reduced['PC2'])

        # Add labels and title
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.title('Scatter Plot')

        # Add labels to each dot
        for i, label in enumerate(embeddings_df.iloc[: , -1].to_list()):
            plt.text(df_reduced['PC1'][i], df_reduced['PC2'][i], label)

        # Save and display the plot
        plt.savefig('../02_Data/principal_component_plot.png', format='png')

    # create and save scatter plot
    create_scatter_plot(df_reduced=df_reduced)

    return df_reduced

# Load embeddings_df.csv into data frame
embeddings_df = pd.read_csv('../02_Data/embeddings_df.csv')

# use the function create_pca_plot to
df_reduced = create_pca_plot(embeddings_df)
01_Data_Pipelines/02_create_PCA_analysis.py
polzerdo55862/vector-store-tutorial
##########################################################################################################
# Calculate cosine similarity between the query vector and all other embedding vectors
##########################################################################################################
import numpy as np
from numpy.linalg import norm
import time
import pandas as pd
import os
import requests

def _get_embeddings(text_chunk):
    '''
    Use embedding model from hugging face to calculate embeddings for the text snippets provided
   
    Parameters:
        - text_chunk (string): the sentence or text snippet you want to translate into embeddings

    Returns:
        - embedding(list): list with all embedding dimensions
    '''
    # define the embedding model you want to use
    model_id = "sentence-transformers/all-MiniLM-L6-v2"

    # you can find the token to the hugging face api in your settings page https://huggingface.co/settings/tokens
    hf_token = os.environ.get('HF_TOKEN')

    # API endpoint for embedding model
    api_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model_id}"
    headers = {"Authorization": f"Bearer {hf_token}"}

    # call API
    response = requests.post(api_url, headers=headers, json={"inputs": text_chunk, "options":{"wait_for_model":True}})
   
    # load response from embedding model into json format
    embedding = response.json()

    return embedding

def calculate_cosine_similarity(text_chunk, embeddings_df):
    '''
    Calculate the cosine similarity between the query sentence and every other sentence
    1. Get the embeddings for the text chunk
    2. Calculate the cosine similarity between the embeddings of our text chunk und every other entry in the data frame

    Parameters:
        - text_chunk (string): the text snippet we want to use to look for similar entries in our database (embeddings_df)
        - embeddings_df (DataFrame): data frame with the columns "text_chunk" and "embeddings"
    Returns:
        -
    '''

    # use the _get_embeddings function the retrieve the embeddings for the text chunk
    sentence_embedding = _get_embeddings(text_chunk)

    # combine all dimensions of the vector embeddings to one array
    embeddings_df['embeddings_array'] = embeddings_df.apply(lambda row: row.values[:-1], axis=1)

    # start the timer
    start_time = time.time()
    print(start_time)

    # create a list to store the calculated cosine similarity
    cos_sim = []

    for index, row in embeddings_df.iterrows():
        A = row.embeddings_array
        B = sentence_embedding

        # calculate the cosine similarity
        cosine = np.dot(A,B)/(norm(A)*norm(B))

        cos_sim.append(cosine)

    embeddings_cosine_df = embeddings_df
    embeddings_cosine_df["cos_sim"] = cos_sim
    embeddings_cosine_df.sort_values(by=["cos_sim"], ascending=False)

    # stop the timer
    end_time = time.time()

    # calculate the time needed to calculate the similarities
    elapsed_time = (end_time - start_time)
    print("Execution Time: ", elapsed_time, "seconds")

    return embeddings_cosine_df

# Load embeddings_df.csv into data frame
embeddings_df = pd.read_csv('../02_Data/embeddings_df.csv')

# test query sentence
text_chunk = "Lilies are white."

# calculate cosine similarity
embeddings_cosine_df = calculate_cosine_similarity(text_chunk, embeddings_df)

# save data frame with text chunks and embeddings to csv
embeddings_cosine_df.to_csv('../02_Data/embeddings_cosine_df.csv', index=False)

# rank based on similarity
similarity_ranked_df = embeddings_cosine_df[["text_chunk", "cos_sim"]].sort_values(by=["cos_sim"], ascending=False)
01_Data_Pipelines/03_calculate_cosine_similarity.py
polzerdo55862/vector-store-tutorial
##########################################################################################################
'''
Includes some functions to create a new vector store collection, fill it and query it
'''
##########################################################################################################
import chromadb
from chromadb.config import Settings
import pandas as pd

# vector store settings
VECTOR_STORE_PATH = r'../02_Data/00_Vector_Store'
COLLECTION_NAME = 'my_collection'

# Load embeddings_df.csv into data frame
embeddings_df = pd.read_csv('../02_Data/embeddings_df.csv')

def get_or_create_client_and_collection(VECTOR_STORE_PATH, COLLECTION_NAME):
    # get/create a chroma client
    chroma_client = chromadb.PersistentClient(path=VECTOR_STORE_PATH)

    # get or create collection
    collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME)

    return collection

# get or create collection
collection = get_or_create_client_and_collection(VECTOR_STORE_PATH, COLLECTION_NAME)

def add_to_collection(embeddings_df):
    # add a sample entry to collection
    # collection.add(
    #     documents=["This is a document", "This is another document"],
    #     metadatas=[{"source": "my_source"}, {"source": "my_source"}],
    #     ids=["id1", "id2"]
    # )

    # combine all dimensions of the vector embeddings to one array
    embeddings_df['embeddings_array'] = embeddings_df.apply(lambda row: row.values[:-1], axis=1)
    embeddings_df['embeddings_array'] = embeddings_df['embeddings_array'].apply(lambda x: x.tolist())

    # add data frame to collection
    collection.add(
        embeddings=embeddings_df.embeddings_array.to_list(),
        documents=embeddings_df.text_chunk.to_list(),
        # create a list of string as index
        ids=list(map(str, embeddings_df.index.tolist()))
    )

# add the embeddings_df to our vector store collection
add_to_collection(embeddings_df)

def get_all_entries(collection):
    # query collection
    existing_docs = pd.DataFrame(collection.get()).rename(columns={0: "ids", 1:"embeddings", 2:"documents", 3:"metadatas"})
    existing_docs.to_excel(r"..//02_Data//01_vector_stores_export.xlsx")
    return existing_docs

# extract all entries in vector store collection
existing_docs = get_all_entries(collection)

def query_vector_database(VECTOR_STORE_PATH, COLLECTION_NAME, query, n=2):
    # query collection
    results = collection.query(
        query_texts=query,
        n_results=n
    )

    print(f"Similarity Search: {n} most similar entries:")
    print(results["documents"])
    return results

# similarity search
similar_vector_entries = query_vector_database(VECTOR_STORE_PATH, COLLECTION_NAME, query=["Lilies are white."])
01_Data_Pipelines/04_set_up_vector_store.py
polzerdo55862/vector-store-tutorial