使用 🤗 transformers、🤗 datasets 和 FAISS 进行多模态数据相似性搜索

嵌入是信息的语义上有意义的压缩。它们可以用于执行相似性搜索、零样本分类或简单地训练新模型。相似性搜索的用例包括在电子商务中搜索类似产品、在社交媒体中搜索内容等等。此笔记本将引导您使用 🤗transformers、🤗datasets 和 FAISS 从特征提取模型创建和索引嵌入，以便以后将其用于相似性搜索。让我们安装必要的库。

!pip install -q datasets faiss-gpu transformers sentencepiece

在本教程中，我们将使用 CLIP 模型来提取特征。CLIP 是一种革命性的模型，它引入了文本编码器和图像编码器的联合训练，以连接两种模态。

import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModel, AutoTokenizer
import faiss
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = AutoModel.from_pretrained("openai/clip-vit-base-patch16").to(device)
processor = AutoImageProcessor.from_pretrained("openai/clip-vit-base-patch16")
tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch16")

加载数据集。为了使此笔记本更轻量级，我们将使用一个小型字幕数据集，jmhessel/newyorker_caption_contest。

from datasets import load_dataset

ds = load_dataset("jmhessel/newyorker_caption_contest", "explanation")

查看一个示例。

>>> ds["train"][0]["image"]

ds["train"][0]["image_description"]

我们不必编写任何函数来嵌入示例或创建索引。 🤗 datasets 库的 FAISS 集成抽象了这些过程。我们可以简单地使用数据集的 map 方法为每个示例创建一列新的嵌入，如下所示。让我们为提示列上的文本特征创建一个。

dataset = ds["train"]
ds_with_embeddings = dataset.map(
    lambda example: {
        "embeddings": model.get_text_features(
            **tokenizer([example["image_description"]], truncation=True, return_tensors="pt").to("cuda")
        )[0]
        .detach()
        .cpu()
        .numpy()
    }
)

我们可以做同样的事情并获取图像嵌入。

ds_with_embeddings = ds_with_embeddings.map(
    lambda example: {
        "image_embeddings": model.get_image_features(**processor([example["image"]], return_tensors="pt").to("cuda"))[
            0
        ]
        .detach()
        .cpu()
        .numpy()
    }
)

现在，我们为每一列创建一个索引。

# create FAISS index for text embeddings
ds_with_embeddings.add_faiss_index(column="embeddings")

# create FAISS index for image embeddings
ds_with_embeddings.add_faiss_index(column="image_embeddings")

使用文本提示查询数据

现在，我们可以使用文本或图像查询数据集以获取其中的相似项。

prmt = "a snowy day"
prmt_embedding = (
    model.get_text_features(**tokenizer([prmt], return_tensors="pt", truncation=True).to("cuda"))[0]
    .detach()
    .cpu()
    .numpy()
)
scores, retrieved_examples = ds_with_embeddings.get_nearest_examples("embeddings", prmt_embedding, k=1)

>>> def downscale_images(image):
...     width = 200
...     ratio = width / float(image.size[0])
...     height = int((float(image.size[1]) * float(ratio)))
...     img = image.resize((width, height), Image.Resampling.LANCZOS)
...     return img


>>> images = [downscale_images(image) for image in retrieved_examples["image"]]
>>> # see the closest text and image
>>> print(retrieved_examples["image_description"])
>>> display(images[0])

['A man is in the snow. A boy with a huge snow shovel is there too. They are outside a house.']

使用图像提示查询数据

图像相似性推理类似，您只需调用 get_image_features。

>>> import requests

>>> # image of a beaver
>>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/beaver.png"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> display(downscale_images(image))

搜索相似图像。

img_embedding = (
    model.get_image_features(**processor([image], return_tensors="pt", truncation=True).to("cuda"))[0]
    .detach()
    .cpu()
    .numpy()
)
scores, retrieved_examples = ds_with_embeddings.get_nearest_examples("image_embeddings", img_embedding, k=1)

显示与海狸图像最相似的图像。

>>> images = [downscale_images(image) for image in retrieved_examples["image"]]
>>> # see the closest text and image
>>> print(retrieved_examples["image_description"])
>>> display(images[0])

['Salmon swim upstream but they see a grizzly bear and are in shock. The bear has a smug look on his face when he sees the salmon.']

保存、推送和加载嵌入

我们可以使用 save_faiss_index 保存带有嵌入的数据集。

ds_with_embeddings.save_faiss_index("embeddings", "embeddings/embeddings.faiss")

ds_with_embeddings.save_faiss_index("image_embeddings", "embeddings/image_embeddings.faiss")

最好将嵌入存储在数据集存储库中，因此我们将创建一个存储库，并将我们的嵌入推送到那里以供以后提取。我们将登录 Hugging Face Hub，在那里创建一个数据集存储库，并将我们的索引推送到那里，并使用 snapshot_download 加载。

from huggingface_hub import HfApi, notebook_login, snapshot_download

notebook_login()

from huggingface_hub import HfApi

api = HfApi()
api.create_repo("merve/faiss_embeddings", repo_type="dataset")
api.upload_folder(
    folder_path="./embeddings",
    repo_id="merve/faiss_embeddings",
    repo_type="dataset",
)

snapshot_download(repo_id="merve/faiss_embeddings", repo_type="dataset", local_dir="downloaded_embeddings")

我们可以使用 load_faiss_index 将嵌入加载到没有嵌入的数据集中。

ds = ds["train"]
ds.load_faiss_index("embeddings", "./downloaded_embeddings/embeddings.faiss")
# infer again
prmt = "people under the rain"

prmt_embedding = (
    model.get_text_features(**tokenizer([prmt], return_tensors="pt", truncation=True).to("cuda"))[0]
    .detach()
    .cpu()
    .numpy()
)

scores, retrieved_examples = ds.get_nearest_examples("embeddings", prmt_embedding, k=1)

>>> display(retrieved_examples["image"][0])

< > Update on GitHub

开源 AI 食谱

使用 🤗 transformers、🤗 datasets 和 FAISS 进行多模态数据相似性搜索

使用文本提示查询数据

使用图像提示查询数据

保存、推送和加载嵌入