Căutare semantică cu FAISS

În secțiunea 5, am creat un dataset cu issues și comentarii din repositoriul 🤗 Datasets. În această secțiune, vom utiliza aceste informații pentru a construi un motor de căutare care ne poate ajuta să găsim răspunsurile la cele mai importante cele mai importante întrebări despre bibliotecă!

Utilizarea embeddings pentru căutare semantică

După cum am văzut în Capitolul 1, Transformer-based language models reprezintă fiecare token într-un fragment de text ca un embedding vector. S-a dovedit că se poate face “pool” embeddingurilor individuale pentru a crea o reprezentare vectorială pentru fraze întregi, paragrafe sau (în anumite cazuri) documente. Aceste embeddings pot fi apoi utilizate pentru a găsi documente similare în corpus prin calcularea dot-product similarity (sau a unei alte metrice de similaritate) între fiecare embedding și returnarea documentelor cu cea mai mare suprapunere.

În această secțiune, vom utiliza embeddings pentru a dezvolta un motor de căutare semantică. Aceste motoare de căutare oferă mai multe avantaje față de abordările convenționale bazate pe căutarea de cuvinte cheie într-un query cu documente.

Încărcarea și pregătirea datasetului

Prima lucru pe care trebuie să îl facem este să descărcăm datasetul nostru cu GitHub issues, așa că folosim funcția load_dataset() ca de obicei:

from datasets import load_dataset

issues_dataset = load_dataset("lewtun/github-issues", 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
})

Aici am specificat splitul default train în load_dataset(), astfel încât returnează un Dataset în loc de DatasetDict. Primul lucru care treubuie făcut este să filtrăm pull requesturile, deoarece acestea rareori tind să fie utilizate pentru a răspunde la întrebările utilizatorilor și vor introduce noise în motorul nostru de căutare. Așa cum ar trebuie deja să știți, putem utiliza funcția Dataset.filter() pentru a exclude aceste rânduri din datasetul nostru. În timp ce suntem aici, putem să filtrăm și rândurile fără comentari, deoarece acestea nu oferă niciun răspuns la întrebările utilizatorilor:

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

Dataset({
    caracteristici: ['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
})

Putem vedea că există multe coloane în datasetul nostru, majoritatea dintre care nu sunt necesare pentru a construi motorul nostru de căutare. Din perspectiva căutării, cele mai informative coloane sunt title, body și comments, în timp ce html_url ne oferă un link înapoi la problema sursă. Hai să utilizăm funcția Dataset.remove_columns() pentru a elimina restul:

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
})

Pentru a crea embeddedurile noastre, vom completa fiecare comentariu cu titlul și body-ul problemei, deoarece aceste câmpuri adesea includ informații contextuale utile. Deoarece coloana noastră comments este în prezent o listă de comentarii pentru fiecare issue, trebuie să “explodăm” coloana, astfel încât fiecare rând să fie format dintr-un tuple (html_url, title, body, comment). În Pandas, putem face acest lucru cu funcția DataFrame.explode(), care creează un rând nou pentru fiecare element dintr-o coloană asemănătoare cu o listă, în timp ce copiază toate celelalte valori ale coloanelor. Pentru a vedea acest lucru în acțiune, să trecem la formatul pandas DataFrame mai întâi:

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

Dacă inspectăm primul rând din acest DataFrame, putem vedea că există patru comentarii asociate acestei probleme:

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...']

Când facem explode df, ne așteptăm să obținem un rând pentru fiecare dintre aceste comentarii. Haideți să verificăm dacă ăsta e cazul:

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

Great, putem vedea că rândurile au fost reproduse, cu coloanele comments incluzând și comentariile individuale. Acum că am terminat cu Pandas, noi putem să schimăm rapid înapoi la un Dataset înărcând DataFrame în memorie:

from datasets import Dataset

comments_dataset = Dataset.from_pandas(comments_df)
comments_dataset

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

Okay, acest lucru ne-a oferit câteva mii de comentarii cu care să lucrăm!

✏️ Încercați! Vezi dacă poți utiliza Dataset.map() pentru a exploda coloana comments din issues_dataset fără a recurge la utilizarea Pandas. Acest lucru este puțin dificil; s-ar putea să găsiți utilă secțiunea “Mapping batch” din documentația 🤗 Datasets pentru această sarcină.

Acum că avem un singur comentariu pe rând, să creăm o nouă coloană comments_length care conține numărul de cuvinte din fiecare comentariu:

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

Putem utiliza această nouă coloană pentru a filtra comentariile scurte, care de obicei includ lucruri precum “cc @lewtun” sau “Mulțumesc!” care nu sunt relevante pentru motorul nostru de căutare. Nu există un număr precis care trebuie selectat pentru filtru, dar aproximativ 15 cuvinte pare a fi un bun punct de plecare:

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
})

După ce am curățat puțin setul nostru de date, putem să concatenăm titlul, descrirea și comentariile problemei împreună într-o nouă coloană text. Ca de obicei, vom scrie o funcție simplă pe care o putem transmite în Dataset.map():

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


comments_dataset = comments_dataset.map(concatenate_text)

Suntem în final pregătiți să creăm niște embeddings! Haideți să vedem cum facem acest lucru.

Crearea embeddings-urilor de text

Am văzut în Capitolul 2 că putem obține token embeddings prin utilizarea clasei AutoModel. Tot ce trebuie să facem este să alegem un checkpoint potrivit pentru a încărca modelul. Din fericire, există o bibliotecă numită sentence-transformers care se ocupă de crearea embeddingurilor. Așa cum se descrie în documentația bibliotecii, cazul nostru este un exemplu de căutare semantică asimetrică deoarece avem o întrebare scurtă al cărei răspuns ne-ar plăcea să îl găsim într-un document mai lung, precum un comentariu la un issue. [Tabelul] (https://www.sbert.net/docs/pretrained_models.html#model-overview) util de prezentare a modelului din documentație indică faptul că checkpointul multi-qa-mpnet-base-dot-v1 are cea mai bună performanță pentru căutarea semantică, așa că îl vom folosi acesta pentru aplicația noastră. De asemenea, vom încărca tokenizer-ul folosind același 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)

Pentru a accelera procesul de embedding, este util să punem modelul și inputurile pe un GPU, deci hai să facem asta acum:

import torch

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

După cum am menționat anterior, dorim să reprezentăm fiecare intrare din corpusul nostru de GitHub issues sub forma unui vector, astfel încât avem nevoie să “agregăm” sau să facem media la tokem embeddings într-un anumit mod. O abordare populară este de a efectua CLS pooling pe outputurile modelului nostru, unde pur și simplu colectăm ultimul stadiu ascuns pentru tokenul special [CLS]. Următoarea funcție face acest lucru pentru noi:

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

În continuare, vom crea o funcție ajutătoare care va tokeniza o listă de documente, va plasa tensorii pe GPU, îi va alimenta în model și, în final, va aplica CLS pooling la outputuri:

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)

Putem testa funcția oferindui prima intrare de text în corpusul nostru și analizând output shapeul:

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

torch.Size([1, 768])

Minunat, am transformat prima intrare din corpusul nostru într-un vector de 768 de dimensiuni! Putem utiliza Dataset.map() pentru a aplica funcția noastră get_embeddings() la fiecare rând din corpusul nostru, astfel încât să creăm o nouă coloană embeddings în acest mod:

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

Observăm că am transformat embeddingurile în matrice NumPy — acest lucru este necesar deoarece biblioteca 🤗 Datasets cere acest format atunci când încercăm să indexăm cu FAISS, ceea ce vom face în continuare.

Utilizarea FAISS pentru căutare de similaritate eficientă

Acum că avem un dataset de embeddings, avem nevoie de o modalitate de a căuta printre ele. Pentru a face acest lucru, vom utiliza o structură de date specială din 🤗 Datasets, numită FAISS index. FAISS (prescurtat de la Facebook AI Similarity Search) este o bibliotecă care oferă algoritmi eficienți pentru căutarea rapidă și clusteringul al embedding vectors.

Ideea de bază din spatele FAISS este crearea unei structuri de date speciale numite index, care permite găsirea embeddingurilor similare cu un input embedding. Crearea unui index FAISS în 🤗 Datasets este simplu — utilizăm funcția Dataset.add_faiss_index() și specificăm care coloană a datasetului nostru dorim să indexăm:

embeddings_dataset.add_faiss_index(column="embeddings")

Acum putem efectua queries pe acest index realizând o căutare a celor mai apropiați vecini cu funcția Dataset.get_nearest_examples(). Haideți să testăm acest lucru prin încorporarea unei întrebări astfel:

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

torch.Size([1, 768])

La fel ca și în cazul documentelor, acum avem un vector de 768 de dimensiuni care reprezintă query-ul, pe care îl putem compara cu întregul corpus pentru a găsi embeddingurile cele mai similare:

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

Funcția Dataset.get_nearest_examples() returnează un tuple cu scorurile care clasifică suprapunerea dintre query și document, și un set corespunzător de sampleuri (în acest caz, cele 5 match-uri). Haideți să colectăm acestea într-un pandas.DataFrame pentru a le putea sorta cu ușurință:

import pandas as pd

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

Acum putem itera peste primele rânduri pentru a veadea cât de bine un query se potrivește cu comentariile disponibile:

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
==================================================
"""

Nu-i rău! A doua încercare se pare că se potrivește cu query-ul!

✏️ Încearcă! Creează propriul tău query și vezi dacp poți găsi un răspuns și să extragi documentele. S-ar putea să trebuiești să crești parametrul k în Dataset.get_nearest_examples() pentru a mări căutarea.