NLP Course documentation

Семантический поиск с помощью FAISS

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Семантический поиск с помощью FAISS

Ask a Question Open In Colab Open In Studio Lab

В разделе 5 мы создали набор данных о issues и комментариях GitHub из репозитория 🤗 Datasets. В этом разделе мы будем использовать эту информацию для создания поисковой системы, которая поможет нам найти ответы на самые насущные вопросы о библиотеке!

Использование эмбеддингов для семанического поиска

Как мы видели в Главе 1, языковые модели на основе Transformer представляют каждую лексему в текстовом фрагменте как эмбеддинг-вектор. Оказывается, можно «объединить» отдельные вложения, чтобы создать векторное представление для целых предложений, абзацев или (в некоторых случаях) документов. Затем эти вложения можно использовать для поиска похожих документов в корпусе путем вычисления скалярного произведения (или какой-либо другой метрики сходства) между каждым вложением и возврата документов с наибольшим перекрытием.

В этом разделе мы будем использовать вложения для разработки семантической поисковой системы. Эти поисковые системы предлагают несколько преимуществ по сравнению с традиционными подходами, основанными на сопоставлении ключевых слов в запросе с документами.

Semantic search.

Загрузка и подготовка датасета

Первое, что нам нужно сделать, это загрузить наш набор данных об issues GitHub, поэтому давайте воспользуемся библиотекой 🤗 Hub для получения URL-адреса, по которому наш файл хранится в Hugging Face Hub:

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",
)

С URL-адресом, сохраненным в data_files, мы можем загрузить удаленный набор данных, используя метод, представленный в раздел 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
})

Здесь мы указали подвыборку train по умолчанию в load_dataset(), поэтому он возвращает Dataset вместо DatasetDict. Первым делом нужно отфильтровать запросы на pull-requests, поскольку они, как правило, редко используются для ответов на вопросы пользователей и создают шум в нашей поисковой системе. Как должно быть уже известно, мы можем использовать функцию Dataset.filter(), чтобы исключить эти строки из нашего набора данных. Давайте также отфильтруем строки без комментариев, поскольку они не дают ответов на запросы пользователей:

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

Мы видим, что в нашем наборе данных много столбцов, большинство из которых нам не нужно для создания нашей поисковой системы. С точки зрения поиска наиболее информативными столбцами являются title, body и comments, а html_url содержит нам ссылку на исходную проблему. Давайте воспользуемся функцией Dataset.remove_columns(), чтобы удалить остальные столбцы:

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

Чтобы создать наши эмбеддинги, мы дополним каждый комментарий заголовком и телом проблемы, поскольку эти поля часто содержат полезную контекстную информацию. Поскольку наш столбец comments в настоящее время представляет собой список комментариев для каждой проблемы, нам нужно «развернуть» столбец, чтобы каждая строка состояла из кортежа (html_url, title, body, comment). В Pandas мы можем сделать это с помощью функции DataFrame.explode(), которая создает новую строку для каждого элемента в столбце, похожем на список, при копировании всех других значений столбца. Чтобы увидеть это в действии, давайте сначала переключимся на формат Pandas DataFrame:

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

Если мы проверим первую строку в этом DataFrame, мы увидим, что есть четыре комментария, связанных с этой проблемой:

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

Когда мы «развернем» df, мы ожидаем получить по одной строке для каждого из этих комментариев. Проверим, так ли это:

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

Отлично, мы видим, что строки были скопированы, а столбец comments содержит отдельные комментарии! Теперь, когда мы закончили с Pandas, мы можем быстро вернуться к Dataset, загрузив DataFrame в память:

from datasets import Dataset

comments_dataset = Dataset.from_pandas(comments_df)
comments_dataset
Dataset({
    features: ['html_url', 'title', 'comments', 'body'],
    num_rows: 2842
})

Хорошо, это дало нам несколько тысяч комментариев для работы!

✏️ Попробуйте! Посмотрите, сможете ли вы использовать Dataset.map(), чтобы развернуть столбец comments столбца issues_dataset без использования Pandas. Это немного сложно; вы можете найти раздел “Batch mapping” документации 🤗 Datasets, полезным для этой задачи.

Теперь, когда у нас есть один комментарий в строке, давайте создадим новый столбец comments_length, содержащий количество слов в комментарии:

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

Мы можем использовать этот новый столбец для фильтрации коротких комментариев, которые обычно содержат такие слова, как «cc @lewtun» или Thanks!», которые не имеют отношения к нашей поисковой системе. Нет точного числа для выбора порога числа слов, но около 15 слов кажется хорошим началом:

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

Немного очистив наш набор данных, давайте соединим название issue, описание и комментарии вместе в новом столбце text. Как обычно, мы напишем простую функцию, которую мы можем передать в Dataset.map():

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


comments_dataset = comments_dataset.map(concatenate_text)

Наконец-то мы готовы создать несколько эмбеддингов’! Давайте взглянем.

Создание текстовых эмбединнгов

В Главе 2 мы видели, что можно получить эмбеддингов токенов с помощью класса AutoModel. Все, что нам нужно сделать, это выбрать подходящую контрольную точку для загрузки модели. К счастью, есть библиотека под названием sentence-transformers, предназначенная для создания эмбеддингов. Как описано в документации библиотеки, наш вариант использования является примером асимметричного семантического поиска потому что у нас есть короткий запрос, ответ на который мы хотели бы найти в более длинном документе, например, в комментарии к проблеме. В удобной таблице обзора модели в документации указано, что контрольная точка multi-qa-mpnet-base-dot-v1 имеет лучшую производительность для семантического поиска, поэтому мы будем использовать её для нашего приложения. Мы также загрузим токенизатор, используя ту же контрольную точку:

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)

Чтобы ускорить процесс построения эмбеддингов, рекомендуется переместить модель и входные данные на устройстве с графическим процессором, поэтому давайте сделаем это сейчас:

import torch

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

Как мы упоминали ранее, мы хотели бы представить каждую запись в нашем корпусе issues GitHub как единый вектор, поэтому нам нужно каким-то образом «объединить» или усреднить наши вложения токенов. Одним из популярных подходов является выполнение CLS pooling выходных данных нашей модели, когда мы просто собираем последнее скрытое состояние для специального токена [CLS]. Следующая функция поможет нам:

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

Далее мы создадим вспомогательную функцию, которая разметит список документов, поместит тензоры в GPU, передаст их в модель и, наконец, применит CLS pooling к выходным данным:

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)

Мы можем проверить работу функции, передав ей первую текстовую запись в нашем корпусе и проверив размерности данных на выходе:

embedding = get_embeddings(comments_dataset["text"][0])
embedding.shape
torch.Size([1, 768])

Отлично, мы преобразовали первую запись в нашем корпусе в 768-мерный вектор! Мы можем использовать Dataset.map(), чтобы применить нашу функцию get_embeddings() к каждой строке в нашем корпусе, поэтому давайте создадим новый столбец embeddings следующим образом:

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

Обратите внимание, что мы преобразовали вложения в массивы NumPy — это потому, что 🤗 Datasets требуют этого формата, когда мы пытаемся проиндексировать их с помощью FAISS, что мы сделаем дальше.

Использование FAISS для эффективного семантического поиска

Теперь, когда у нас есть датасет с эмбеддингами, нам нужен способ поиска по ним. Для этого мы будем использовать специальную структуру данных из 🤗 Datasets, называемую FAISS index. FAISS (сокращение от Facebook AI Similarity Search) — это библиотека, предоставляющая эффективные алгоритмы для быстрого поиска и кластеризации эмбеддингов.

Основная идея FAISS состоит в том, чтобы создать специальную структуру данных, называемую index, которая позволяет найти, какие эмбеддинги подобны входным эмбеддингам. Создать индекс FAISS в 🤗 Datasets очень просто — мы используем функцию Dataset.add_faiss_index() и указываем, какой столбец нашего набора данных мы хотим проиндексировать:

embeddings_dataset.add_faiss_index(column="embeddings")

Теперь мы можем выполнять запросы к этому индексу, выполняя поиск ближайшего соседа с помощью функции Dataset.get_nearest_examples(). Давайте проверим это, сначала внедрив вопрос следующим образом:

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

Как и в случае с документами, теперь у нас есть 768-мерный вектор, представляющий запрос, который мы можем сравнить со всем корпусом, чтобы найти наиболее похожие объекты:

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

Функция Dataset.get_nearest_examples() возвращает набор оценок, которые ранжируют совпадение между запросом и документом, и соответствующий набор образцов (здесь 5 лучших совпадений). Давайте соберем их в pandas.DataFrame, чтобы мы могли легко их отсортировать:

import pandas as pd

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

Теперь мы можем пройтись по первым нескольким строкам, чтобы увидеть, насколько наш запрос соответствует имеющимся комментариям:

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

Неплохо! Наше второе обращение, кажется, соответствует запросу.

✏️ Попробуйте! Создайте свой собственный запрос и посмотрите, сможете ли вы найти ответ в найденных документах. Возможно, вам придется увеличить параметр k в Dataset.get_nearest_examples(), чтобы расширить поиск.