Ricerche vettoriali e modelli di embedding con Chroma DB e OpenAI

La ricerca vettoriale nei testi consente di trovare informazioni a partire da relazioni semantiche e concetti complessi. Questo metodo può essere impiegato per creare motori di ricerca avanzati che integrano le conoscenze e la comprensione del linguaggio dei Large Language Models (LLMs).

Ti è mai capitato di effettuare una ricerca in un database ma senza sapere quali parole chiave utilizzare? Probabilmente sapevi come descrivere ciò che ti serviva, ma senza i termini giusti i risultati non erano pertinenti come speravi.
Le ricerche tradizionali sono infatti incentrate sull’utilizzo di parole chiave e i risultati restituiti sono basati soprattutto su criteri come la corrispondenza tra termini e la frequenza con cui determinate parole si presentano in un insieme di testi. 
In questo caso, la qualità e la pertinenza dei risultati dipende in maniera importante dall’utilizzo esplicito di determinati termini nella nostra query di ricerca.

Come funziona la ricerca vettoriale?

La ricerca vettoriale si avvale di algoritmi di calcolo della distanza tra vettori numerici, usati per rappresentare il significato di parole e frasi, al fine di restituire i risultati semanticamente più vicini a una descrizione testuale fornita come input.
In altre parole, la ricerca vettoriale ti aiuta a trovare quello che cerchi in un mare di testi non strutturati, anche se non conosci le parole chiave esatte che dovresti usare. Per questo motivo è un concetto così importante nell’ambito dei Large Language Models (LLMs).

Ma cosa sono esattamente questi vettori? Sono essenzialmente rappresentazioni numeriche dei testi, dette vector embeddings, che catturano il significato e le relazioni tra le parole. I numeri all'interno di questi vettori sono generati tramite modelli di machine learning e rappresentano specifiche caratteristiche o attributi del testo.
Più due testi sono simili tra loro, più i numeri nei rispettivi vettori di embeddings saranno collocati in prossimità nello spazio vettoriale. Ciò significa che la distanza dei vettori tra due testi che parlano di “gatti” e “felini” sarà minore rispetto a quella tra due testi che parlano di “pesci” e “montagne”.
Con questo meccanismo, la ricerca vettoriale permette quindi di effettuare ricerche più avanzate, individuando relazioni complesse che portano a risultati più ricchi e contestualizzati.

• "Un film di fantascienza con viaggi nel tempo e paradossi temporali"
• "Un dramma storico sulla lotta per i diritti civili"
• "Un film d'azione ambientato durante una pandemia globale"

Tramite il meccanismo della ricerca vettoriale, le descrizioni fornite vengono dunque tradotte in vettori numerici, i quali sono poi comparati con i vettori delle trame dei film presenti nel database.
Ciò porta a una selezione di film più pertinente e coerente con il genere che realmente stai cercando. 

Se questo esempio della ricerca del film ideale ti ha intrigato, buone notizie!
In questo articolo lavoreremo con il dataset IMDB Movies Dataset - Top 1000 Movies by IMDB Rating per configurare un database vettoriale con Chroma DB. Ci serviremo inoltre delle API di OpenAI per la generazione degli embeddings a partire dai titoli e dalle descrizioni dei film.

Il risultato finale consisterà in un piccolo motore di ricerca semantico che ci permetterà di trovare film di nostro interesse utilizzando descrizioni testuali delle trame, senza dover dipendere da categorie predefinite.

Requisiti tecnici

• Python 3.10 per scrivere il codice sorgente
• Chroma DB, un database vettoriale open-source, per memorizzare i vettori numerici e per eseguire le query di ricerca
• OpenAI Embeddings per generare i vettori basati sulle descrizioni dei film.

Prima di procedere con lo sviluppo, eseguiamo il seguente comando nel terminale per installare le librerie necessarie.

pip install pandas, chromadb

Setup Chroma DB e OpenAI Embeddings API

Creiamo ora il file chroma_search.py nella directory del progetto. Importiamo le librerie necessarie, includiamo la nostra chiave API di OpenAI e assegniamo il modello "text-embedding-3-small" a una nuova variabile. Questo modello ci consentirà di generare i vettori necessari.

Inizializziamo in seguito il client chromadb, che creerà il percorso del nostro database vettoriale nella directory di progetto. Successivamente, creiamo la collezione "imdb_movies", che conterrà tutti i vettori generati dalle descrizioni dei film che verranno fornite in seguito.

import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
import pandas as pd

# Insert your OpenAI API key
OPENAI_API_KEY = "sk-XXXXXXXXX"

# Define the embedding model to be used for text embedding
EMBEDDING_MODEL = "text-embedding-3-small"

# Initialize an instance of OpenAIEmbeddingFunction with the API key and embedding model
openai_embedding_function = OpenAIEmbeddingFunction(api_key=OPENAI_API_KEY, model_name=EMBEDDING_MODEL)

# Initialize a persistent ChromaDB client
client_chromadb = chromadb.PersistentClient(path="chromadb")

# Get or create the "imdb_movies" collection in ChromaDB
collection = client_chromadb.get_or_create_collection("imdb_movies", embedding_function=openai_embedding_function)

Import e preparazione del dataset

In questa fase, importeremo il nostro dataset e trasformeremo i dati al suo interno per darli in pasto al modello di embedding. Per fare ciò, creeremo tre nuove funzioni.

La funzione import_movie_dataset importa un dataset CSV e restituisce un dataframe Pandas.
Poiché i film nel file originale non avevano un identificatore univoco, useremo l'indice del dataframe come ID incrementale. Questo passaggio è cruciale in quanto Chroma DB richiede esplicitamente un ID per ogni documento nella sua raccolta, senza il quale l'operazione di inserimento dati non andrebbe a buon fine.
In questa fase, convertiamo inoltre l'anno di uscita del film in formato numerico. Questo ci consentirà di filtrare i risultati in base all'anno, ad esempio mostrando tutti i film usciti dopo il 2010.

def import_movie_dataset(path: str = "dataset/imdb_top_1000.csv") -> pd.DataFrame:
    """
    Import movie dataset from the specified path and return it as a pandas DataFrame.

    Args:
        path (str, optional): The file path to the dataset. Defaults to "dataset/imdb_top_1000.csv".

    Returns:
        pd.DataFrame: A DataFrame containing the imported movie dataset.
    """
    movies = pd.read_csv(path)
    movies["id"] = (movies.index + 1).astype(str)

    # Convert released year to integer
    movies["Released_Year"] = movies["Released_Year"].str.replace('[^0-9]', '0', regex=True).fillna(0).astype(int)

    return movies

La funzione create_text_from_movie ha il compito di assemblare il corpus di testo che verrà utilizzato per generare i vettori mediante l'embedding del modello.
I contenuti selezionati in questa fase determineranno i valori all’interno dei nostri vettori e influenzeranno la pertinenza dei risultati di ricerca. Poiché le possibili query di ricerca saranno principalmente descrizioni generiche di film, abbiamo deciso di creare il corpus di testo includendo titolo, genere e una breve descrizione del film.

def create_text_from_movie(movie:dict) -> str:
    """
    Returns a string containing the title, genre, and description of the movie
    """
    return f"""Title: {movie["Series_Title"]}
    Genre: {movie["Genre"]}
    Description: {movie["Overview"]}
    """

La funzione extract_metadata_from_movie è opzionale e prevede l'estrazione dei metadati dei film per l'inserimento in Chroma DB, oltre al corpus di testo dei documenti. L’aggiunta di questi metadati può essere utile se desideri filtrare le tue ricerche in base a parametri extra, come l'anno di uscita dei film che vuoi includere nei risultati.

def extract_metadata_from_movie(movie: dict, meta_keys: list[str]) -> dict:
    """
    Returns a dictionary with metadata from the movie based on the selected keys.
    """
    metadata = {}
    for key in meta_keys:
        if key in movie and pd.notnull(movie[key]):
            metadata[key] = movie[key]
    return metadata

Scrittura dei dati nel database

La funzione add_movie_vectors_to_db provvede alla scrittura dei dati nel database vettoriale, utilizzando il metodo upsert del package chromadb. Questo metodo è idempotente, il che significa che crea un nuovo documento se non esiste già e lo aggiorna se è già presente nel database con lo stesso ID.
Poiché si tratta di un'operazione dispendiosa in termini di prestazioni, si consiglia di eseguirla soltanto al momento della scrittura iniziale dei dati o in caso di futuri aggiornamenti al database, in modo da evitare esecuzioni ridondanti ad ogni utilizzo del motore di ricerca.

def add_movie_vectors_to_db(
    movies: pd.DataFrame,
    collection: chromadb.Collection,
    meta_keys: list[str] = None
) -> None:
    """
    Adds movie vectors to the specified ChromaDB collection.

    Args:
    - movies: DataFrame containing movie data.
    - collection: Collection in the database to which movie vectors will be added.
    - meta_keys: List of metadata keys to extract from movies. Defaults to None.
    """

    # Convert DataFrame to dictionary
    movies_dict = movies.to_dict(orient="records")

    # Create text embeddings for each movie in the dataset
    movies_text = [create_text_from_movie(movie) for movie in movies_dict]

    # Prepare list of movie IDs
    movies_ids = [movie["id"] for movie in movies_dict]

    if meta_keys:
        # Select and extract metadata fields for each movie
        movies_meta = [extract_metadata_from_movie(movie, meta_keys) for movie in movies_dict]

        try: 
            # Upsert movie vectors and metadata into the database collection
            collection.upsert(
                ids=movies_ids,
                documents=movies_text,
                metadatas=movies_meta
            )
            print("UPLOADED MOVIES TO DB WITH METADATA")

        except:
            print("FAILED TO LOAD MOVIES TO DB")
    else:
        try: 
            # Upsert movie vectors into the database collection without metadata
            collection.upsert(
                ids=movies_ids,
                documents=movies_text
            )
            print("UPLOADED MOVIES TO DB WITHOUT METADATA")

        except:
            print("FAILED TO LOAD MOVIES TO DB")

Per scrivere i dati nella collection di Chroma DB possiamo eseguire i seguenti comandi dal nostro terminale python. Il terzo argomento della funzione add_movie_vectors_to_db è opzionale e permette di utilizzare i valori nelle colonne selezionate come metadati dei documenti.

>>> from chroma_search import client_chromadb, \
						   collection, \
                           import_movie_dataset, \
                           create_text_from_movie, \
                           extract_metadata_from_movie, \
                           add_movie_vectors_to_db
>>> movies = import_movie_dataset()
>>> add_movie_vectors_to_db(movies, collection, ["Released_Year", "IMDB_Rating", "Director", "Gross"])

Utilizzo delle query di ricerca

Siamo quasi pronti per effettuare le nostre prime query sulla collection di Chroma DB per ottenere subito risultati di film. Per poterlo fare, dichiariamo prima la funzione query_text_vector_db.

def query_text_vector_db(
    collection: chromadb.Collection,
    query_text: str,
    dataframe: pd.DataFrame,
    n_results: int = 5,
    where_clause: dict = {}
) -> pd.DataFrame:
    """
    Perform a vector search on a chromadb collection from a query text and returns a DataFrame with results sorted by the nearest vector distances.

    Args:
        collection (chromadb.Collection): A collection object representing the chromadb collection to be searched.
        query_text (str): A query text to search for in the collection.
        dataframe (pd.DataFrame): A DataFrame containing the data associated with the collection.
        n_results (int, optional): Number of results to retrieve for each query text. Defaults to 5.
        where_clause (dict, optional): Additional query constraints. Defaults to {}.

    Returns:
        pd.DataFrame: A DataFrame containing the search results
    """
    
    # Query the collection with provided query texts
    results = collection.query(
            query_texts=query_text,
            n_results=n_results,
            where=where_clause
        )
    
    # Get result ids for the query text
    term_result_ids = results["ids"][0]

    # Filter the main DataFrame to get the suggested movies
    suggested_movies = dataframe.copy()[dataframe["id"].isin(term_result_ids)]
    
    # Add a column with the vector distance of each result from the query vector
    suggested_movies["vector_distance"] = results["distances"][0]

    # Sort search results by vector distance for relevance
    suggested_movies = suggested_movies.sort_values(by=["vector_distance"])

    # Define columns to filter from the DataFrame
    filter_cols = [
        'vector_distance',
        'Series_Title',
        'Overview',
        'Released_Year',
        'Director'
    ]
    
    return suggested_movies[filter_cols]

La funzione di ricerca restituisce un dataframe Pandas con i risultati ordinati per rilevanza e accetta come argomenti:

• la collezione chromadb con i nostri film
• una query di ricerca
• il numero di risultati desiderati
• il dataframe originario con i film da estrarre
• dei filtri aggiuntivi opzionali basati sui metadati

Analizziamo ora i punti principali che compongono la nostra funzione.

Struttura dei risultati di ricerca

Per comprendere la struttura di risposta salvata nella variabile intermedia results, si riporta come esempio il risultato della query "a movie about wizards", con il numero di risultati limitato eccezionalmente a uno per facilitarne la lettura. 

{
  "ids": [["782"]],
  "distances": [[0.28548556566238403]],
  "metadatas": [[
    {
      "Director": "Mike Newell",
      "Gross": "290,013,036",
      "IMDB_Rating": 7.7,
      "Released_Year": 2005
    }
  ]],
  "embeddings": null,
  "documents": [
    [
      "Title: Harry Potter and the Goblet of Fire\n    Genre: Adventure, Family, Fantasy\n    Description: Harry Potter finds himself competing in a hazardous tournament between rival schools of magic, but he is distracted by recurring nightmares.\n    "
    ]
  ],
  "uris": null,
  "data": null
}

Possiamo osservare che di default la chiave "embeddings" non restituisce il valore del vettore dei risultati, poiché è molto lungo e poco interessante per la lettura dei risultati. Al contrario, un valore significativo è rappresentato dalla chiave "distances", che indica la distanza (calcolata sulla base della similarità del coseno) tra il vettore generato dalla nostra query di testo e il vettore associato al film estratto, in questo caso "Harry Potter and the Goblet of Fire".

In sostanza, la risposta fornisce i primi n risultati ordinati in base alla vicinanza tra i loro vettori e il vettore della query inserita.

È interessante notare che, nonostante il testo di descrizione del film utilizzato per alimentare il motore di ricerca non contenga affatto la parola "wizards", il modello di embeddings utilizzato per generare i vettori ha comunque catturato le relazioni semantiche tra le parole, consentendo così di individuare un risultato altamente pertinente e significativo.

Utilizzo dei filtri sui metadati

Grazie all'aggiunta dei metadati ai documenti in Chroma DB, abbiamo ora la possibilità di servirci anche del parametro where_clause nella funzione di ricerca. Ad esempio, se volessimo cercare solo i film usciti prima del 2000, potremmo passare il seguente dizionario come valore.

where_clause = {
    "Released_Year": {
        "$lt": 2000
    }
}

I filtri di ricerca possono inoltre essere concatenati tra loro tramite gli operatori AND e OR. Per un elenco esaustivo si rimanda alla documentazione di Chroma DB relativa all'uso dei filtri.

Presentazione dei risultati

Di seguito è riportata la funzione main del nostro file. Possiamo eseguirla con una query a nostro piacimento per interrogare il database usando la ricerca vettoriale.

def main(query):
    # Importing the movie dataset
    movies = import_movie_dataset()

    # Querying the text vector database to find relevant search results
    search_results = query_text_vector_db(collection=collection, query_text=query, dataframe=movies)

    print(search_results)
    
if __name__ == "__main__":
    main("A cyberpunk movie")

Ecco qui i risultati ottenuti tramite la query "A cyberpunk movie".

Limitazioni

Il nostro motore di ricerca non stabilisce un limite minimo di somiglianza tra i vettori per restituire i risultati. Ciò significa che se non ci sono elementi rilevanti per la query di ricerca, verranno comunque mostrati i risultati che sono matematicamente più vicini, anche se non pertinenti.
Per valutare l'accuratezza dei risultati, possiamo esaminare i valori nella chiave "distances" e stabilire una soglia di accettazione per determinare quali risultati sono da considerare rilevanti.

Per fini didattici, abbiamo utilizzato un dataset con pochi record e descrizioni brevi, contenenti al massimo una o due frasi per trama di film. Questo significa che il nostro motore di ricerca dispone di una conoscenza limitata dei dettagli delle trame dei film e può essere ampiamente influenzato da descrizioni o titoli vaghi dei film.
Per esempio, la query "a movie about car racing" pone al terzo posto il film "Gran Torino", nonostante siano presenti altri film più pertinenti alla query in questione.
Il motore di ricerca è in questo caso influenzato dal fatto che il titolo del film contiene il nome di un'auto, sebbene la pellicola non tratti di corse automobilistiche.

Nella progettazione del nostro motore di ricerca vettoriale, è quindi fondamentale concentrarsi sulla scelta e sulla preparazione accurata dei dati testuali, assicurandoci di disporre di una mole di dati sufficiente per garantire la pertinenza dei risultati.

Conclusioni

In questo articolo siamo partiti da una spiegazione teorica delle ricerche vettoriali, descrivendo come funziona il processo di recupero delle informazioni attraverso i vector embeddings, che sono essenzialmente rappresentazioni numeriche dei dati testuali in uno spazio vettoriale.

Sulla base di queste nozioni, abbiamo creato un motore di ricerca semantico utilizzando Chroma DB, un database vettoriale open-source. Questo strumento consente di agevolare la ricerca del film ideale, sfruttando la comprensione del linguaggio fornita dagli embeddings di OpenAI.

Note aggiuntive su quanto sviluppato:

• Le tecniche descritte in questo articolo possono essere applicate in modo simile per creare sistemi di raccomandazione che si basano sullo stesso concetto di vicinanza dei vettori. Ad esempio, è possibile utilizzare il vettore di un film presente nel nostro database come punto di partenza per cercare altri film simili
• Oltre a Chroma DB, ci sono molte altre opzioni da considerare per progettare il proprio database vettoriale, tra cui alcune soluzioni open-source e altre commerciali. Lo stesso vale per i modelli di embedding diversi da quelli di OpenAI, come ad esempio quelli disponibili su Hugging Face, che fornisce una vasta gamma di modelli open-source
• L’utilizzo di una maggior quantità di dati testuali da elaborare in forma vettoriale può portare a risultati di ricerca più precisi e contestualizzati. Tuttavia, è importante considerare le limitazioni sul numero massimo di token che i modelli possono gestire e adottare strategie di chunking del contenuto, al fine ottimizzare la rilevanza dei risultati delle nostre ricerche.