Sviluppare un motore di ricerca vettoriale con Langchain e Streamlit

In questo progetto sfrutteremo le abilità linguistiche dei Large Language Models (LLMs) per lo sviluppo di un motore di ricerca vettoriale. Questo motore ci consentirà di effettuare ricerche dettagliate tra una grande quantità di dati testuali non strutturati.  

In un articolo precedente, abbiamo introdotto il concetto e il funzionamento delle ricerche vettoriali con dati testuali. In sintesi, questo metodo di ricerca si basa su algoritmi di calcolo della distanza tra vettori numerici, utilizzati per rappresentare semanticamente parole e frasi, al fine di restituire i risultati più pertinenti e simili al significato dei testi utilizzati come query. 
A differenza delle ricerche tradizionali che si basano sulla corrispondenza tra parole chiave, la ricerca vettoriale tiene conto anche delle relazioni concettuali tra la query di ricerca e i testi memorizzati in un database vettoriale, garantendo risultati più contestualizzati e accurati.

Applicazioni pratiche delle ricerche vettoriali

In ambito aziendale capita spesso di confrontarsi con una grande quantità di dati testuali non organizzati e non sempre ben categorizzati. All'interno di questi dati possono tuttavia nascondersi informazioni preziose alle quali vorremmo poter accedere in modo rapido ed efficiente.

Ecco alcuni esempi di applicazioni pratiche della ricerca vettoriale in diversi contesti:

• E-commerce: Permettere all'utente di effettuare ricerche più descrittive e articolate per trovare il prodotto giusto da acquistare
• Raccomandazione di contenuti: Fornire agli utenti suggerimenti su contenuti simili da visualizzare, in base alla somiglianza semantica dei testi consultati in precedenza
• Supporto clienti: Offrire ai clienti un sistema che li agevoli nella ricerca di risposte a domande e dubbi a partire da un'ampia raccolta di domande frequenti (FAQ)
• Knowledge base interna: Mettere a disposizione del team un sistema di ricerca della documentazione interna per rispondere alle specifiche esigenze operative

Prendendo spunto dall'esempio proposto nell'articolo introduttivo, svilupperemo un motore di ricerca vettoriale che ci permetterà di trovare il film ideale in base alla descrizione della trama. Il dataset che utilizzeremo è Wikipedia Movie Plots with AI Plot Summaries, che include oltre 34'000 trame di pellicole cinematografiche, sia in versione completa che riassunta.

Ecco di seguito il risultato che otterremo al termine del progetto.

Perché usare Langchain?

Langchain è ormai diventato il framework di riferimento per lo sviluppo di applicazioni basate sui Large Language Models (LLMs), grazie anche alla facilità con cui consente di integrare tra loro componenti utili e strumenti di terze parti.
Il suo approccio consente di passare facilmente da una libreria esterna all'altra mantenendo buona parte della struttura di codice invariata. Questa flessibilità ci permette quindi di sperimentare facilmente con l'utilizzo di diversi modelli, per comparare i risultati ottenuti e ottimizzare la nostra soluzione.

Nel nostro progetto, sfrutteremo questa caratteristica di Langchain per permettere all'utente di scegliere il modello da utilizzare per le ricerche. Le due opzioni a disposizione sono date dai modelli 'all-MiniLM-L6-v2', un sentence transformer disponibile gratuitamente su Hugging Face, e 'text-embedding-3-small' di OpenAI.

Requisiti tecnici

• Python 3.10 per scrivere il codice sorgente
• Langchain come framework di riferimento per integrare e orchestrare le componenti del motore di ricerca
• Chroma DB, un database vettoriale open-source, per memorizzare i vettori numerici e per eseguire le query di ricerca
• OpenAI Embeddings e Sentence Transformers per generare i vettori basati sulle descrizioni dei film
• Streamlit per la creazione di un'interfaccia grafica user-friendly tramite la quale interrogare il database vettoriale.

Prima di procedere con lo sviluppo, salviamo il file csv contenente i film nella cartella dataset ed eseguiamo il seguente comando nel terminale per installare le librerie necessarie.

pip install python-dotenv pandas chromadb langchain langchain_openai sentence-transformers streamlit

Al termine dell'installazione delle libreria possiamo creare il file app.py e importare i package necessari. È inoltre necessario creare la variabile d'ambiente nel file OPENAI_API_KEY con la chiave API di OpenAI all'interno del file .env.

import pandas as pd
from langchain_community.document_loaders import DataFrameLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma
import streamlit as st
import os
from dotenv import load_dotenv

load_dotenv()

Preparazione dei contenuti testuali

Prima di addentrarci negli aspetti tecnici della preparazione dei dati, è essenziale comprendere una peculiarità dei vettori numerici con i quali lavoreremo. Ciascun modello di embeddings genera un vettore di dimensione fissa, a prescindere dalla lunghezza del testo di input.
Per esempio, il modello 'text-embedding-3-small' produce vettori di 1'536 numeri, che rappresentano specifici attributi e caratteristiche del testo. Questo significa che il vettore generato da una trama di film composta da 10 paragrafi avrà la stessa lunghezza di quello per una trama composta da una sola frase.

Nonostante la dimensione dei vettori che andremo a generare con ciascun modello resterà sempre invariata, la lunghezza dei testi forniti per la loro generazione ricopre comunque un ruolo di rilievo in fase di preparazione del contenuto. È infatti importante considerare due aspetti:

1. I vari modelli di embeddings presentano ciascuno diverse lunghezze ottimali di caratteri in input per garantire le migliori performance e catturare al meglio l'essenza dei testi
2. Ogni modello di embeddings prevede una limitazione al numero massimo di caratteri (context window) che possono essere processati o convertiti in vettore in una singola operazione.

Per far fronte a questi vincoli e limitazioni, si adotta solitamente una pratica nota come "chunking", che consiste nel suddividere il testo in blocchi più piccoli di una lunghezza predefinita. Questo permette ai modelli di embeddings di elaborare ogni chunk di testo nel rispetto della loro context window e di massimizzare la capacità di catturare relazioni semantiche, per costruire rappresentazioni vettoriali significative.
Questo approccio è fondamentale nel contesto di sviluppo di applicativi LLMs, come ad esempio i chatbot basati sul meccanismo Retrieval-Augmented Generation (RAG), i quali sono incentrati su un sistema mirato di domande e risposte che prevede il reperimento di informazioni granulari all'interno dei testi.

Nel caso del nostro progetto, questa pratica rischia di introdurre un effetto indesiderato, in quanto i film con trame più lunghe finiranno per generare un maggior numero di chunk che saranno indicizzati nel database. Da ciò ne consegue una maggior probabilità che questi film appaiano nei risultati di ricerca, a scapito di altri film altrettanto rilevanti ma con trame più brevi.

Una strategia alternativa sarebbe quella di servirsi di un LLM per riassumere il testo di nostro interesse, al fine di produrre trame di lunghezza uniforme e idonea per i modelli di embedding che utilizzeremo.

Nel caso del nostro dataset, questo procedimento è già stato svolto e potremo dunque utilizzare il contenuto nella colonna 'PlotSummary', riassunto a partire dalla colonna 'Plot'.

La funzione load_csv_to_docs carica il file CSV in formato dataframe Pandas tramite il loader di Langchain, al fine di formattarne il contenuto in documenti che potranno essere inseriti nel database vettoriale.
L'argomento content_col della funzione ci permette di specificare da quale colonna del dataset estrarre il testo per creare gli embeddings, mentre le colonne restanti del dataframe verranno invece considerate come metadati per i documenti.

def load_csv_to_docs(file_path:str="./dataset/wiki_movie_plots_deduped_with_summaries.csv", 
                     content_col:str="PlotSummary"
                     ) -> list:
    """
    Load a CSV file into documents using Langchain DataFrame loader.

    Args:
        file_path (str): The file path to the CSV file.
        content_col (str): The name of the column containing the content of each document.

    Returns:
        list: A list of documents loaded from the CSV file.
    """

    df = pd.read_csv(file_path)

    loader = DataFrameLoader(df, page_content_column=content_col)

    documents = loader.load()

    return documents

La funzione successiva, split_docs_to_chunks, prende in input i documenti generati dalla funzione precedente e li suddivide in chunk testuali pronti per essere elaborati dal modello di embedding.

def split_docs_to_chunks(documents:list, chunk_size:int=1000, chunk_overlap:int=0) -> list:
    """
    Split documents into chunks and format each chunk.

    Args:
        documents (list): A list of documents to be split.
        chunk_size (int, optional): The size of each chunk. Defaults to 1000.
        chunk_overlap (int, optional): The overlap between consecutive chunks. Defaults to 0.

    Returns:
        list: A list of formatted chunks.
    """
    # Create a RecursiveCharacterTextSplitter instance
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    # Split documents into chunks using the text splitter
    chunks = text_splitter.split_documents(documents)
    
    # Iterate over each chunk
    for chunk in chunks:
        # Extract metadata from the chunk
        title = chunk.metadata['Title']
        origin = chunk.metadata['Origin/Ethnicity']
        genre = chunk.metadata['Genre']
        release_year = chunk.metadata['Release Year']
        
        # Extract content from the chunk
        content = chunk.page_content
        
        # Format the content with metadata
        final_content = f"TITLE: {title}\nORIGIN: {origin}\nGENRE: {genre}\nYEAR: {release_year}\nBODY: {content}\n"
        
        # Update the page content of the chunk with formatted content
        chunk.page_content = final_content
    
    return chunks

Questo passaggio ha due scopi principali:

1. Garantire che ogni chunk non superi una determinata lunghezza, in questo caso impostata su un default di 1'000 caratteri, per evitare che il testo ecceda la context window del modello. Nel nostro caso lavoreremo con contenuti la cui dimensione è già al di sotto di questa soglia, ma questa funzione può essere utile da includere comunque nel programma, nel caso in cui si desiderasse sperimentare con l'utilizzo di testi alternativi che superano la lunghezza massima.
2. Fornire ai chunk informazioni contestuali aggiuntive, come i metadati dei film, che saranno incluse nella stringa di testo utilizzata per la generazione dei vector embeddings. Questa pratica migliora l'accuratezza delle operazioni di recupero dei vettori, grazie al contesto globale introdotto nei chunk.

Configurazione di Chroma DB

La funzione create_or_get_vectorstore sfrutta le due funzioni già create in precedenza per creare una nuova istanza di database vettoriale Chroma DB, o caricarne una se già esistente nella directory del progetto. Questa funzione gestisce anche il processo di generazione dei vettori a partire dal nostro dataset, basandosi sul modello di embedding scelto come argomento. Le opzioni disponibili per il modello di embedding sono 'OpenAI' e 'SentenceTransformer'.

def create_or_get_vectorstore(file_path: str, content_col: str, selected_embedding: str) -> Chroma:
    """
    Create or get a Chroma vector store based on the selected embedding model.

    Args:
        file_path (str): The file path to the dataset.
        content_col (str): The name of the column containing the content of each document.
        selected_embedding (str): The selected embedding model ('OpenAI' or 'SentenceTransformer').

    Returns:
        Chroma: A Chroma vector store.
    """
    # Determine the embedding function and database path based on the selected embedding model
    if selected_embedding == 'OpenAI':
        embedding_function = OpenAIEmbeddings(model="text-embedding-3-small", chunk_size=100, show_progress_bar=True)
        db_path = './chroma_openai'

    elif selected_embedding == 'SentenceTransformer':
        embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
        db_path = './chroma_hf'

    # Check if the database directory exists
    if not os.path.exists(db_path):
        # If the directory does not exist, create the database
        documents = load_csv_to_docs(file_path, content_col)
        chunks = split_docs_to_chunks(documents)

        print("CREATING DB...")
        db = Chroma.from_documents(chunks, embedding_function, persist_directory=db_path)

    else:
        # If the directory exists, load the existing database
        print('LOADING DB...')
        db = Chroma(persist_directory=db_path, embedding_function=embedding_function)

    return db

È importante notare che a seconda del modello di embedding scelto, verranno create due istanze di database separate in percorsi diversi. I vettori generati dai due modelli avranno lunghezze e parametri diversi, quindi è essenziale in seguito caricare le rispettive istanze di database utilizzando la stessa chiave di embedding che è stata usata per la loro creazione.

Nel codice per l'embedding di OpenAI, abbiamo specificato il valore del parametro chunk_size con un valore di 100. Questo parametro, al contrario di quanto può far credere il nome, non applica un chunking effettivo al contenuto, ma determina quanti documenti processare al massimo in ciascun batch per generare gli embedding. È consigliabile specificare questo parametro con un valore inferiore rispetto al suo default di 1'000, per evitare errori di limite di utilizzo delle API di OpenAI. Impostando inoltre show_progress_bar su True, visualizziamo una barra di avanzamento durante il processo di embedding, che dovrebbe richiedere circa 10 minuti per il nostro dataset.

Per l'embedding con il Sentence Transformer, non è invece necessaria la suddivisione in batch e il processo di vettorializzazione dovrebbe completarsi nel giro di pochi secondi.

Ora possiamo procedere creando due istanze del nostro database vettoriale eseguendo i seguenti comandi nel terminale Python.

>>> from app import load_csv_to_docs, \
				    split_docs_to_chunks, \
                    create_or_get_vectorstore
>>> documents = load_csv_to_docs()
>>> chunks = split_docs_to_chunks(documents)
>>> create_or_get_vectorstore('./dataset/wiki_movie_plots_deduped_with_summaries.csv','PlotSummary','OpenAI')
>>> create_or_get_vectorstore('./dataset/wiki_movie_plots_deduped_with_summaries.csv','PlotSummary','SentenceTransformer')

Gestione delle query di ricerca

La nostra ultima funzione è chiamata query_vectostore e ha lo scopo di trovare i primi k risultati più rilevanti tramite ricerca vettoriale nel database.

def query_vectorstore(db:Chroma, 
                      query:str, 
                      k:int=20, 
                      filter_dict:dict={}
                      ) -> pd.DataFrame:
    """
    Query a Chroma vector store for similar documents based on a query.

    Args:
        db (Chroma): The Chroma vector store to query.
        query (str): The query string.
        k (int, optional): The number of similar documents to retrieve. Defaults to 20.
        filter_dict (dict, optional): A dictionary specifying additional filters. Defaults to {}.

    Returns:
        pd.DataFrame: A DataFrame containing metadata of the similar documents.
    """
    # Perform similarity search on the vector store
    results = db.similarity_search(query, filter=filter_dict, k=k)

    # Initialize an empty list to store metadata from search results
    results_metadata = []

    # Extract metadata from results
    for doc in results:
        results_metadata.append(doc.metadata)

    # Convert metadata to a DataFrame
    df = pd.DataFrame(results_metadata)
    
    # Drop duplicate rows based on the 'Wiki Page' column
    df.drop_duplicates(subset=['Wiki Page'], keep='first', inplace=True)

    return df

Questa funzione richiede come input obbligatori l'istanza del database vettoriale da interrogare e una query di ricerca testuale. Inoltre, è possibile specificare il numero di risultati desiderati (il default è 20) e applicare filtri di ricerca ai metadati tramite un dizionario. Nella nostra interfaccia, includeremo infatti un box che permetterà all'utente di filtrare i risultati in base all'anno di uscita del film.

È importante notare che ogni chunk creato in precedenza corrisponde a una diversa entità di documento nel nostro database. Questo significa che una ricerca può restituire diversi documenti riferiti allo stesso film. Per garantire che i film proposti non siano duplicati, abbiamo incluso nella funzione un passaggio finale di deduplicazione dei risultati. Nel nostro caso, abbiamo comunque creato un solo chunk per ogni film, ma se si lavora con testi più lunghi, è importante tenere conto del fatto che il numero di risultati unici potrebbe essere inferiore a quello specificato nel parametro k della funzione.

Sviluppo della UI con Streamlit

Ora che disponiamo di tutte le funzioni necessarie per far funzionare il motore di ricerca, possiamo concentrarci sullo sviluppo di un'interfaccia che permetta all'utente di interagire in modo semplice con l'applicazione.

Per costruire questa interfaccia useremo Streamlit, un framework gratuito e open-source che facilita la trasformazione di uno script Python in un'applicazione web.
La logica del funzionamento di Streamlit comporta l'esecuzione completa del nostro script Python ogni volta che l'utente interagisce con l'interfaccia, ad esempio inserendo un nuovo input in un modulo.

Procediamo inserendo l'intera logica dell'applicazione Streamlit nella funzione main del nostro file.

La prima parte di questa funzione si occupa di gestire il toggle per la scelta del modello da utilizzare per le ricerche vettoriali. A seconda del modello selezionato, viene inizializzata l'istanza del database contenente i vettori corrispondenti.

def main():
	# Apply Streamlit page config
    st.set_page_config(
    page_title=" Vector Search Engine | Datasense",
    page_icon="https://143998935.fs1.hubspotusercontent-eu1.net/hubfs/143998935/Datasense_Favicon-2.svg"
    )

    # Read and apply custom CSS style
    with open('./css/style.css') as f:
        css = f.read()
    st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)
    
    # Display logo and title
    st.image("https://143998935.fs1.hubspotusercontent-eu1.net/hubfs/143998935/Datasense%20Logo_White.svg", width=180)
    st.title("Building a Vector Search Engine with Langchain and Streamlit")
    st.markdown("Find your ideal movie among over 30k films with an AI-powered vector search engine.")

    # Toggle for using OpenAI embeddings
    if "openai_on" not in st.session_state:
        st.session_state.openai_on = False

    openai_on = st.toggle('Use OpenAI embeddings')

    # Check if the toggle value has changed
    if openai_on != st.session_state.openai_on:
        # Clear the existing database from the session state
        if "db" in st.session_state:
            del st.session_state.db

    # Determine selected embedding model
    if openai_on:
        selected_embedding = "OpenAI"
        st.session_state.openai_on = True
    else:
        selected_embedding = "SentenceTransformer"
        st.session_state.openai_on = False

    # Create or get the vector store database
    file_path = './dataset/wiki_movie_plots_deduped_with_summaries.csv'
    content_col = 'PlotSummary'
    st.session_state.db = create_or_get_vectorstore(file_path, content_col, selected_embedding)
    
    # ...

La parte successiva della funzione predispone l'area di testo in cui l'utente può inserire la sua query di ricerca e gestisce la creazione di un dizionario per filtrare la ricerca, opzionalmente passato alla funzione query_vectorstore.

Questo filtro consente di selezionare i film usciti prima, dopo o in un anno specifico. I filtri attivi vengono visualizzati in una sezione apposita assegnata alla variabile filter_box.

def main():

    # ...
    
    # Text input for query
    query = st.text_input("Tonight I'd like to watch...", "A thriller movie about memory loss")

    # Display filter options
    filter_box = st.empty()
    with st.expander("Filter by year"):
        with st.form("filters"):
            filter_dict = {}
            st.write("Release year...")
            year_operator = st.selectbox(
                label="Operator",
                options=("is equal to", "is after", "is before")
            )
            year = st.number_input(
                label="Year",
                min_value=1900,
                max_value=2023,
                value=2000
            )
            submitted = st.form_submit_button("Apply filter")
            operator_signs = {
                "is equal to": "$eq",
                "is after": "$gt",
                "is before": "$lt"
            }

            if submitted:
                filter_dict = {
                    "Release Year": {
                        f"{operator_signs[year_operator]}": year
                    }
                }
                # Escape the HTML tags
                filter_box.markdown(
                    f"<p><b>Active filter</b>:</p> <span class='active-filter'>Released year {year_operator} {year}</span>", 
                    unsafe_allow_html=True
                )
                
     # ...

Nella fase finale, la funzione esegue la ricerca nel database vettoriale utilizzando la funzione query_vectorstore. Qui, forniamo il database vettoriale basato sul modello prescelto, la query testuale e, se desiderato dall'utente, dei filtri di ricerca aggiuntivi.

Dato che questa funzione restituisce un dataframe, i risultati della ricerca vengono visualizzati iterando su ogni riga. Le colonne contenenti i metadati pertinenti vengono quindi selezionate e presentate in un riquadro dedicato al film.

def main():
    # ...

    # Perform search if query exists
    if query:
        # Perform vector store query
        results_df = query_vectorstore(
            db=st.session_state.db,
            query=query,
            filter_dict=filter_dict
        )

        # Display search results
        for index, row in results_df.iterrows():
            # Escape the HTML tags
            st.markdown(
                f"""
                <div class='result-item-box'>
                    <span class='label-genre'{row['Genre']}</span>
                    <h4>{row['Title']}</h4>
                    <div class='metadata'>
                        <p><b>Year:</b> {row['Release Year']}</p>
                        <p><b>Director:</b> {row['Director']}</p>
                        <p><b>Origin:</b> {row['Origin/Ethnicity']}</p>
                    </div>
                    <a href='{row['Wiki Page']}'>Read more →</a>
                </div>
                """,
                unsafe_allow_html=True
            )

if __name__ == '__main__':
    main()

Per avviare il nostro motore di ricerca basterà eseguire il seguente comando nel terminale.

streamlit run app.py

L'applicazione si aprirà automaticamente sulla porta 8501 in localhost.

Conclusione

In questo progetto, abbiamo visto come la ricerca vettoriale possa offrire un modo efficace e contestualizzato per accedere a grandi quantità di dati testuali non strutturati, in particolare se combinata con la conoscenza del linguaggio dei LLMs e i relativi vettori di embedding.

Per la creazione del nostro motore di ricerca, abbiamo utilizzato Langchain, un framework che semplifica l'integrazione di varie componenti legate ai LLMs, facilitando così la sperimentazione tra modelli e l'ottimizzazione delle prestazioni dell'applicazione.

Durante la preparazione dei dati testuali, abbiamo evidenziato l'importanza della qualità dei dati forniti ai modelli di embedding. In particolare, abbiamo mostrato come il chunking e/o il riassunto dei contenuti possano essere utilizzati per ottimizzare i dati testuali prima di procedere con la loro vettorializzazione.

Il risultato finale del nostro progetto è un motore di ricerca vettoriale dotato di un'interfaccia basata su Streamlit, che permette agli utenti di cercare tra un'ampia raccolta di film tramite due diversi modelli di embedding.

Se sei interessato alla codebase completa del motore di ricerca, puoi trovarla nella repository Github del progetto.