Ricerca semantica con FAISS

Nella sezione 5 abbiamo creato un dataset di issue e commenti dalla repository GitHub di 🤗 Datasets. In questa sezione useremo queste informazioni per costrure un motore di ricerca semantico che ci può aiutare a trovare risposte alle nostre domande urgenti sulla libreria!

Usare gli embedding per la ricerca semantica

Come abbiamo visto nel Capitolo 1, i language model basati su Transformer rappresentano ogni token in un testo come un vettore, detto embedding. È possibile “mettere insieme” i diversi embedding per creare una rappresentazione vettoriale di un’intera frase, paragrafo o (in alcuni casi) documento. Questi embedding possono essere usati per trovare documenti simili in un corpus calcolandone la similarità, ad esempio usando il prodotto scalere (o altre misure di similarità) tra ogni embedding, e restituendo i documenti più simili.

In questa sezione useremo gli embedding per sviluppare un motore di ricerca semantico. Questi motori di ricerca offrono diversi vantagig rispetto ai metodo convenzionali, basati sulla ricerca, all’interno dei documenti, delle parole chiavi presente in una query.

Caricare e preparare il dataset

La prima cosa che dobbiamo fare è scaricare il nostro dataset di issue, quindi utilizziamo la libreria 🤗 Hub per scaricare i file usando l’URL dell’Hub Hugging Face:

from huggingface_hub import hf_hub_url

data_files = hf_hub_url(
    repo_id="lewtun/github-issues",
    filename="datasets-issues-with-comments.jsonl",
    repo_type="dataset",
)

Se conseriamo l’URL iin data_files, possiamo caricare il dataset utilizzando il metodo introdotto nella sezione 2:

from datasets import load_dataset

issues_dataset = load_dataset("json", data_files=data_files, split="train")
issues_dataset

Dataset({
    features: ['url', 'repository_url', 'labels_url', 'comments_url', 'events_url', 'html_url', 'id', 'node_id', 'number', 'title', 'user', 'labels', 'state', 'locked', 'assignee', 'assignees', 'milestone', 'comments', 'created_at', 'updated_at', 'closed_at', 'author_association', 'active_lock_reason', 'pull_request', 'body', 'performed_via_github_app', 'is_pull_request'],
    num_rows: 2855
})

Qui abbiamo specificato la sezione di defaul train in load_dataset(), così che questa funzione resituisce un Dataset invece di un DatasetDict. La prima cosa da fare è filtrare le richieste di pull, poichè queste tendono a essere usate raramente come risposta alle domande degli utenti, e introdurrebbero rumore nel nostro motore di ricerca. Come dovrebbe esser enoto, possiamo usare la funzione Dataset.filter() per escludere questi dati dal nostro dataset. Già che ci siamo, eliminiamo anche le righe senza commenti, poiché queste non presentano nessuna risposta alle domande degli utenti:

issues_dataset = issues_dataset.filter(
    lambda x: (x["is_pull_request"] == False and len(x["comments"]) > 0)
)
issues_dataset

Dataset({
    features: ['url', 'repository_url', 'labels_url', 'comments_url', 'events_url', 'html_url', 'id', 'node_id', 'number', 'title', 'user', 'labels', 'state', 'locked', 'assignee', 'assignees', 'milestone', 'comments', 'created_at', 'updated_at', 'closed_at', 'author_association', 'active_lock_reason', 'pull_request', 'body', 'performed_via_github_app', 'is_pull_request'],
    num_rows: 771
})

Possiamo vedere che ci sono molte colonne nel nostro dataset, molte delle quali non servono alla costruzione del nostro motore di ricerca. Da una prospettiva di ricerca, le colonne maggiormente informative sono title, body, e comments, mentre html_url ci fornisce un link all’issue originale. Usiamo la funzione Dataset.remove_columns() per eliminare le colonne rimanenti:

columns = issues_dataset.column_names
columns_to_keep = ["title", "body", "html_url", "comments"]
columns_to_remove = set(columns_to_keep).symmetric_difference(columns)
issues_dataset = issues_dataset.remove_columns(columns_to_remove)
issues_dataset

Dataset({
    features: ['html_url', 'title', 'comments', 'body'],
    num_rows: 771
})

Per crare i nostri embedding arricchiremo ognu commento con il titolo e il corpo dell’issue, visto che questi campi spesso includono informazioni utili sul contesto. Poiché la nostra colonna comment è al momento una lista di commenti per ogni issue, dobbiamo “farla esplodere” così che ogni riga consista in una tupla (html_url, title, body, comment). In panda è possibile farlo utilizzando la funzione Dataframe.explode(), che crea una nuova riga per ogni elemento in una colonna in formato di lista, ripetendo i valori di tutte le altre colonne. Per vederlo in azione, prima di tutto passiamo al formato DataFrame:

issues_dataset.set_format("pandas")
df = issues_dataset[:]

Se diamo un’occhiata alla prima riga di questo DataFrame, possiamo vedere che ci sono quattro commenti associati con quest’issue:

df["comments"][0].tolist()

['the bug code locate in ：\r\n    if data_args.task_name is not None:\r\n        # Downloading and loading a dataset from the hub.\r\n        datasets = load_dataset("glue", data_args.task_name, cache_dir=model_args.cache_dir)',
 'Hi @jinec,\r\n\r\nFrom time to time we get this kind of `ConnectionError` coming from the github.com website: https://raw.githubusercontent.com\r\n\r\nNormally, it should work if you wait a little and then retry.\r\n\r\nCould you please confirm if the problem persists?',
 'cannot connect，even by Web browser，please check that  there is some  problems。',
 'I can access https://raw.githubusercontent.com/huggingface/datasets/1.7.0/datasets/glue/glue.py without problem...']

Quando “esplodiamo” df, ci aspettiamo di avere una riga per ognuno di questi commenti. Controlliamo se è così:

comments_df = df.explode("comments", ignore_index=True)
comments_df.head(4)

	html_url	title	comments	body
0	https://github.com/huggingface/datasets/issues/2787	ConnectionError: Couldn't reach https://raw.githubusercontent.com	the bug code locate in ：\r\n if data_args.task_name is not None...	Hello,\r\nI am trying to run run_glue.py and it gives me this error...
1	https://github.com/huggingface/datasets/issues/2787	ConnectionError: Couldn't reach https://raw.githubusercontent.com	Hi @jinec,\r\n\r\nFrom time to time we get this kind of `ConnectionError` coming from the github.com website: https://raw.githubusercontent.com...	Hello,\r\nI am trying to run run_glue.py and it gives me this error...
2	https://github.com/huggingface/datasets/issues/2787	ConnectionError: Couldn't reach https://raw.githubusercontent.com	cannot connect，even by Web browser，please check that there is some problems。	Hello,\r\nI am trying to run run_glue.py and it gives me this error...
3	https://github.com/huggingface/datasets/issues/2787	ConnectionError: Couldn't reach https://raw.githubusercontent.com	I can access https://raw.githubusercontent.com/huggingface/datasets/1.7.0/datasets/glue/glue.py without problem...	Hello,\r\nI am trying to run run_glue.py and it gives me this error...

bene, possiamo vedere che le righe sono state duplicate, e che la colonna comment contiene i diversi comment! Ora che abbiamo finito con Pandas, possiamo passare velocemente a Dataset caricando il DataFrame in memoria:

from datasets import Dataset

comments_dataset = Dataset.from_pandas(comments_df)
comments_dataset

Dataset({
    features: ['html_url', 'title', 'comments', 'body'],
    num_rows: 2842
})

Perfetto, ora abbiamo qualche migliaio di commenti con cui lavorare!

✏️ Prova tu! Prova ad utilizzare Dataset.map() per far esplodere la colonna commenti di issues_dataset senza utilizzare Pandas. È un po’ difficile: potrebbe tornarti utile la sezione “Batch mapping” della documentazione di 🤗 Datasets.

Ora che abbiamo un commento per riga, creiamo una nuova colonna comments_length che contiene il numero di parole per ogni commento:

comments_dataset = comments_dataset.map(
    lambda x: {"comment_length": len(x["comments"].split())}
)

Possiamo usare questa nuova colonna per eliminare i commenti brevi, che solitamente includono cose del tipo “cc @lewtun” o “Grazie!”, che non sono pertinenti per il nostro motore di ricerca. Non abbiamo un numero preciso da selezionare per questo filtro, ma 15 parole dovrebbero andare bene:

comments_dataset = comments_dataset.filter(lambda x: x["comment_length"] > 15)
comments_dataset

Dataset({
    features: ['html_url', 'title', 'comments', 'body', 'comment_length'],
    num_rows: 2098
})

Una volta data una pulizia al nostro dataset, possiamo concatenare il titolo, la descrizione e i commenti delle issue in una nuova colonna text. Come al solito , scriveremo una semplice funzione che possiamo passare a Dataset.map():

def concatenate_text(examples):
    return {
        "text": examples["title"]
        + " \n "
        + examples["body"]
        + " \n "
        + examples["comments"]
    }


comments_dataset = comments_dataset.map(concatenate_text)

Siamo finalmente pronti a creare degli embedding! Diamo un’occhiata.

Creare i text embedding

Abbiamo visto nel Capitolo 2 che possiamo ottenere i token embedding utilizando la classe AutoModel. Dobbiamo solo scegliere un checkpoint valido da cui caricare il modell. Per fortuna, esiste una libreria chiamata sentence-transformers, dedicata alla creazione di embedding. Seguendo la descrizione nella documentazionedella libreria, il nostro caso d’uso è un esempio di asymmetric semantic search perché abbiamo una breve query per cui vogliamo trovare risposte in un documento lungo, come ad esempio un commento a un issue. La scheda di riepilogo dei modelli nella documentazione ci indica che il checkpoint multi-qa-mpnet-base-dot-v1 ha mostrato la performance migliore per la ricerca semantica, quindi è quello che useremo per la nostra applicazione. Caricheremo anche il tokenizzatore usando lo stesso checkpoint:

from transformers import AutoTokenizer, AutoModel

model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)

Per accelerare il processo di embedding, è bene usare la GPU per il modello e gli input, quindi:

import torch

device = torch.device("cuda")
model.to(device)

Come abbiamo già detto prima, vorremmo rappresentare ogni entrata nel nostro corpus di issue GitHub come un vettore singolo, per cui avremo bisogno di calcolare la media, o il “pool” dei nostri token embedding. Un metodo comune è di effettuare un CLS pooling sull’output del nostro modello: questa tecnica su basa sul recuperare semplicemente l’ultimo stato nascosto del token speciale [CLS]. La funzione seguente fa proprio questo:

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]

Poi, creeremo una funzione di supporto che: tokenizza una lista di documenti, inserire i tensori sulla GPU, li usa come input per il modello, e infine applica il CLS pooling agli output:

def get_embeddings(text_list):
    encoded_input = tokenizer(
        text_list, padding=True, truncation=True, return_tensors="pt"
    )
    encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
    model_output = model(**encoded_input)
    return cls_pooling(model_output)

Possiamo testare la funzione sul primo testo nel nostro corpus, e ispezionandone le dimensioni dell’ouput:

embedding = get_embeddings(comments_dataset["text"][0])
embedding.shape

torch.Size([1, 768])

Bene, abbiamo convertito la prima voce del nostro corpus in un vettore a 768 dimensioni! Possiamo usare Dataset.map() per applicare la nostra funzione get_embedding() a ogni riga del nostro corpus, quindi creiamo una nuova colonna embedding così:

embeddings_dataset = comments_dataset.map(
    lambda x: {"embeddings": get_embeddings(x["text"]).detach().cpu().numpy()[0]}
)

Node che abbiamo convertito gli embedding in array NumPy — questo perchè 🤗 Datasets ha bisogno di questo formato per indicizzare gli embedding con FAISS, che è ciò che faremo nella prossima sezione.

Usare FAISS per ricerca di similarità efficiente

Ora che abbiamo un dataset di embedding, abbiamo bisogno di un modo per effettuare una ricerca. Per far ciò, useremo una struttura specialie di 🤗 Datasets chiamato indice FAISS. FAISS (Facebook AI Similarity Search) è una libreria che permette di utilizzare algoritmi efficient per ricercare e raggruppare gli embedding.

L’idea di base dietro FAISS è di creare un formato speciale di dati chiamato indice che permette di trovare quali embedding sono simili a un embedding in input. Creare un indice FAISS su 🤗 Datasets è semplice — usiamo la funzione Dataset.add_faiss_index() e specificare quale colonna nel nostro dataset vorremmo indicizzare:

embeddings_dataset.add_faiss_index(column="embeddings")

Ora possiamo eseguire dele query su questo indice effettuando una ricerca degli elementi più vicini usando la funzione Dataset.get_nearest_examples(). Testiamolo creando un embedding per una domanda.

question = "How can I load a dataset offline?"
question_embedding = get_embeddings([question]).cpu().detach().numpy()
question_embedding.shape

torch.Size([1, 768])

Proprio come con i documenti, ora abbiamo un vettore di 768 dimensioni che rappresenta la query, che possiamo confrontare con l’intero corpus per trovare gli embedding più simili:

scores, samples = embeddings_dataset.get_nearest_examples(
    "embeddings", question_embedding, k=5
)

La funzione Dataset.get_nearest_examples() restituisce una tupla di valori che valutano la sovrapposizione tra la query e il documento, e un set corrispondente di campioni (in questo caso, le 5 corrispondenze migliori). Salviamole in un pandas.DataFrame, così che possiamo ordinarle facilmente:

import pandas as pd

samples_df = pd.DataFrame.from_dict(samples)
samples_df["scores"] = scores
samples_df.sort_values("scores", ascending=False, inplace=True)

Ora possiamo iterare sulle prime righe per vedere quanto bene la nostra query corrisponde ai commenti disponibili:

for _, row in samples_df.iterrows():
    print(f"COMMENT: {row.comments}")
    print(f"SCORE: {row.scores}")
    print(f"TITLE: {row.title}")
    print(f"URL: {row.html_url}")
    print("=" * 50)
    print()

"""
COMMENT: Requiring online connection is a deal breaker in some cases unfortunately so it'd be great if offline mode is added similar to how `transformers` loads models offline fine.

@mandubian's second bullet point suggests that there's a workaround allowing you to use your offline (custom?) dataset with `datasets`. Could you please elaborate on how that should look like?
SCORE: 25.505046844482422
TITLE: Discussion using datasets in offline mode
URL: https://github.com/huggingface/datasets/issues/824
==================================================

COMMENT: The local dataset builders (csv, text , json and pandas) are now part of the `datasets` package since #1726 :)
You can now use them offline
\`\`\`python
datasets = load_dataset("text", data_files=data_files)
\`\`\`

We'll do a new release soon
SCORE: 24.555509567260742
TITLE: Discussion using datasets in offline mode
URL: https://github.com/huggingface/datasets/issues/824
==================================================

COMMENT: I opened a PR that allows to reload modules that have already been loaded once even if there's no internet.

Let me know if you know other ways that can make the offline mode experience better. I'd be happy to add them :)

I already note the "freeze" modules option, to prevent local modules updates. It would be a cool feature.

----------

> @mandubian's second bullet point suggests that there's a workaround allowing you to use your offline (custom?) dataset with `datasets`. Could you please elaborate on how that should look like?

Indeed `load_dataset` allows to load remote dataset script (squad, glue, etc.) but also you own local ones.
For example if you have a dataset script at `./my_dataset/my_dataset.py` then you can do
\`\`\`python
load_dataset("./my_dataset")
\`\`\`
and the dataset script will generate your dataset once and for all.

----------

About I'm looking into having `csv`, `json`, `text`, `pandas` dataset builders already included in the `datasets` package, so that they are available offline by default, as opposed to the other datasets that require the script to be downloaded.
cf #1724
SCORE: 24.14896583557129
TITLE: Discussion using datasets in offline mode
URL: https://github.com/huggingface/datasets/issues/824
==================================================

COMMENT: > here is my way to load a dataset offline, but it **requires** an online machine
>
> 1. (online machine)
>
> ```
>
> import datasets
>
> data = datasets.load_dataset(...)
>
> data.save_to_disk(/YOUR/DATASET/DIR)
>
> ```
>
> 2. copy the dir from online to the offline machine
>
> 3. (offline machine)
>
> ```
>
> import datasets
>
> data = datasets.load_from_disk(/SAVED/DATA/DIR)
>
> ```
>
>
>
> HTH.


SCORE: 22.893993377685547
TITLE: Discussion using datasets in offline mode
URL: https://github.com/huggingface/datasets/issues/824
==================================================

COMMENT: here is my way to load a dataset offline, but it **requires** an online machine
1. (online machine)
\`\`\`
import datasets
data = datasets.load_dataset(...)
data.save_to_disk(/YOUR/DATASET/DIR)
\`\`\`
2. copy the dir from online to the offline machine
3. (offline machine)
\`\`\`
import datasets
data = datasets.load_from_disk(/SAVED/DATA/DIR)
\`\`\`

HTH.
SCORE: 22.406635284423828
TITLE: Discussion using datasets in offline mode
URL: https://github.com/huggingface/datasets/issues/824
==================================================
"""

Non male! Il nostro secondo risultato sembra soddisfare la nostra richiesta.

✏️ Prova tu! Crea la tua query e prova a trovare una risposta tra i documenti raccolti. Potresti aver bisogno di aumentare il parametro k in Dataset.get_nearest_examples() per allargare la ricerca.

NLP Course

Ricerca semantica con FAISS

Usare gli embedding per la ricerca semantica

Caricare e preparare il dataset

Creare i text embedding

Usare FAISS per ricerca di similarità efficiente