使用BGE-M3筛选图片

发表于 分类于 阅读次数： Waline：

最近在做多模态LLM的业务，需要从开源数据集中筛选一些符合我需求的图片，想到了使用 BGE-M3 结合图像 caption 数据来实现这个功能。

准备 sharegpt4v 数据

首先准备 sharegpt4v 的数据，在sharegpt4v有详细说明，在此不再赘述。

使用 bge-m3 进行向量检索

在这里直接给代码吧，应该很容易看懂，document 直接使用图像的 caption 数据，因此不用再使用 langchain 中的各种 loader。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from langchain_core.documents.base import Document
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

# 如果没有 gpu 的话，'device'给'cpu'
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-m3", model_kwargs = {'device': 'cuda'})

with open("share-captioner_coco_lcs_sam_1246k_1107.json", 'r') as file:
    captioner_json = json.load(file.read())

docs = []
for doc_json in captioner_json:
    curr_doc = Document(doc_json['conversations'][1]['value'])
    curr_doc.metadata['id'] = doc_json['id']
    curr_doc.metadata['image_path'] = doc_json['image']

    docs.append(curr_doc)

db = FAISS.from_documents(docs, embeddings)
search_docs = db.similarity_search("some landscape photos", k=10)

因为在 document 的 meta 信息中添加了图片的 id 和路径，因此可以很容易地在向量检索完成之后，找到原始的图片链接。