bert+np.memap+faiss文本相似度匹配 topN

目录

任务

代码

结果说明

---

任务

使用 bert-base-chinese 预训练模型将文本数据向量化后，使用 np.memap 进行保存，再使用 faiss 进行相似度匹配出每个文本与它最相似的 topN

此篇文章使用了地址数据，目的是为了跑通这个流程，数据可以自己构建

模型下载：bert预训练模型下载-CSDN博客

np.memap ：

是NumPy库中的一种内存映射文件（Memory-Mapped File）对象，它允许你将硬盘上的大文件以类似数组的方式访问和操作，而不需要一次性将整个文件加载到内存中。

当你创建一个numpy.memmap对象时，实际上是创建了一个与磁盘文件对应的虚拟数组。当读取或修改这个数组中的元素时，NumPy会自动在需要的时候从磁盘上读取或写入相应的数据块。这样做的好处在于：

1. 节省内存：对于超过系统可用内存的大文件，可以通过内存映射文件进行处理。
2. 高效性：由于操作系统对内存映射文件的支持，频繁的小规模I/O操作可以得到优化。

在NLP领域，内存映射文件常用于存储大量的文本向量化结果，如BERT嵌入向量等

faiss：

是由Facebook AI Research团队开发的一款开源库，专注于大规模相似性搜索和聚类问题。它特别适用于稠密向量的高效索引和检索，为高维向量空间中的近似最近邻（Approximate Nearest Neighbor, ANN）搜索提供了强大的支持。

主要特点：

1. 高性能：Faiss在CPU和GPU上都具有出色的性能表现，能够处理十亿级别的向量数据集，并且提供了一系列优化过的索引结构和算法以适应不同的应用场景。

2. 多种索引方法：包括基于树的索引、量化索引（如IVF、PQ等）、层次化索引以及基于图的索引方法，这些索引方式允许在精度与速度之间进行权衡。

3. 内存和磁盘持久化：支持将索引保存到磁盘并在需要时加载，这对于大型或持久化应用至关重要。

4. 多平台兼容：Faiss由C++编写，但提供了Python接口，同时也支持其他语言的绑定。

5. 灵活的应用场景：广泛应用于图像搜索、推荐系统、信息检索、文本分类和聚类等多个领域，其中涉及的向量通常是通过深度学习模型生成的嵌入向量。

6. 易于使用：API设计简洁明了，开发者可以快速地构建高效的相似度搜索系统。

代码

先对文本数据进行向量化，串行的，to_vector.py

import torch
import numpy as np
import pandas as pd
import time
from transformers import BertTokenizer, BertModel
from tqdm.auto import tqdm


class TextEmbedder():
    def __init__(self, model_name="../bert-base-chinese"):
        # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 自己电脑跑不起来 gpu
        self.device = torch.device("cpu")
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = BertModel.from_pretrained(model_name).to(self.device)

    def embed_sentences(self, sentences):
        inputs = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt').to(self.device)
        # 创建一个进度条，用于显示句子列表的处理进度
        sentences_iterator = enumerate(inputs['input_ids'])
        progress_bar = tqdm(sentences_iterator, total=len(inputs['input_ids']), desc="Embedding Progress")

        embeddings_list = []
        for i, input_id in progress_bar:
            with torch.no_grad():
                # 对每个句子单独获取嵌入向量
                embedding = self.model(input_ids=input_id.unsqueeze(0)).last_hidden_state[:, 0, :].cpu().numpy()
                embeddings_list.append(embedding)

        # 将所有句子的嵌入向量堆叠成最终的嵌入矩阵
        embeddings = np.vstack(embeddings_list)
        return embeddings

    def save_embeddings_to_memmap(self, sentences, output_file, dtype=np.float32):
        embeddings = self.embed_sentences(sentences)
        shape = embeddings.shape
        embeddings_memmap = np.memmap(output_file, dtype=dtype, mode='w+', shape=shape)
        embeddings_memmap[:] = embeddings[:]
        del embeddings_memmap  # 关闭并确保数据已写入磁盘


def read_data():
    data = pd.read_excel('../data/data.xlsx')
    return data['address'].to_list()


def main():
    # text_data = ["这是第一个句子", "这是第二个句子", "这是第三个句子"]
    text_data = read_data()

    embedder = TextEmbedder()

    # 设置输出文件路径
    output_filepath = '../output/123sentence_embeddings.npy'

    # 将文本数据向量化并保存到内存映射文件
    embedder.save_embeddings_to_memmap(text_data, output_filepath)


if __name__ == "__main__":
    start = time.time()
    main()
    end = time.time()
    print(end - start)

在上面代码中，embedding 是串行计算的，如果用了 gpu 不能充分利用资源，所以要改成批次方式，单独修改 embed_sentences 方法：

class TextEmbedder():
    # ...（其余类定义不变）

    def embed_sentences(self, sentences, batch_size=32):
        inputs = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt').to(self.device)
        
        # 计算批次数量
        batch_count = (len(inputs['input_ids']) + batch_size - 1) // batch_size

        embeddings_list = []
        with tqdm(total=len(sentences), desc="Embedding Progress") as pbar:
            for batch_idx in range(batch_count):
                start = batch_idx * batch_size
                end = min((batch_idx + 1) * batch_size, len(inputs['input_ids']))

                # 提取当前批次的输入ID
                current_batch_input_ids = inputs['input_ids'][start:end]

                with torch.no_grad():
                    # 对当前批次获取嵌入向量
                    embedding_batch = self.model(current_batch_input_ids).last_hidden_state[:, 0, :].cpu().numpy()

                # 将当前批次的嵌入向量添加到列表中
                embeddings_list.extend(embedding_batch.tolist())
                
                # 更新进度条
                pbar.update(end - start)

        # 将所有批次的嵌入向量堆叠成最终的嵌入矩阵
        embeddings = np.vstack(embeddings_list)
        return embeddings

再向量化后会生成一个保存的向量文件，使用向量文件找出每个文本与它最相似的 topN，find_topN.py

import pandas as pd
import numpy as np
import faiss
from tqdm import tqdm


def search_top4_similarities(index_path, data):
    index = faiss.IndexFlatL2(768)  # 假设BERT输出维度是768
    embeddings_memmap = np.memmap(index_path, dtype=np.float32, mode='r')

    # 确保embeddings_memmap是二维数组，如有需要转换
    if len(embeddings_memmap.shape) == 1:
        embeddings_memmap = embeddings_memmap.reshape(-1, 768)

    index.add(embeddings_memmap)

    results = []
    for i, text_emb in enumerate(tqdm(embeddings_memmap)):
        D, I = index_nature.search(np.expand_dims(text_emb, axis=0), topk)  # 查找前topk个最近邻
        
        # 获取对应的 nature_df_img_id 的索引
        top_k_indices = I[0][:10]  # topk 10
        # 根据索引提取 nature_df_img_id
        top_k_ids = [nature_df_img_id[index] for index in top_k_indices]

        # 计算余弦相似度并构建字典
        cosine_similarities = [cosine_similarity(text_emb, embeddings_memmap[index]) for index in top_k_indices]
        top_similarity = dict(zip(top_k_ids, cosine_similarities))

        results.append((data['shop_id'].to_list()[i], top_similarity))

    return results


# 使用余弦相似度公式，这里假设 cosine_similarity 是一个计算两个向量之间余弦相似度的函数
def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))


def main_search():
    data = pd.read_excel('../data/data.xlsx')
    similarities = search_top4_similarities('../output/sentence_embeddings.npy', data)

    # 输出结果
    similar_df = pd.DataFrame(similarities, columns=['id', 'top4'])
    similar_df.to_csv('../output/similarities.csv', index=False)

# 执行搜索并保存结果
main_search()

结果说明

输出格式有两列如下：

id top4 id {id1：相似度, id2：相似度, id3：相似度， id4：相似度}

以上代码流程跑通目的达到了，以后有更长的文本内容就可以使用了

这里去核对结果，发现效果并不是很好，可能原因地理数据还是太短了，分词效果可能不一样

1. BERT模型的局限性：尽管BERT是一个强大的预训练语言模型，但它在处理地址这类具有地理信息和特定结构的数据时可能不是最佳选择。因为BERT是在大量自然语言文本上预训练的，对于地址这种格式化的数据，它可能无法捕捉到地名之间精确的地理位置关系。

2. 相似度计算方法：在上面的示例代码中，我们使用的是欧氏距离作为相似度的代理，对于文本向量而言，这可能不是最佳的选择，尤其是在BERT这样的Transformer模型输出的嵌入空间中，余弦相似度通常更适用于衡量文本向量间的相似性。

3. 数据预处理：地址标准化是关键步骤，如果地址没有经过充分的预处理（如去掉无关词、统一行政区划名称等），那么即使是相似的地址也可能得到不同的向量化表示。

4. 索引构建和搜索参数：Faiss库提供了多种索引方式，针对不同规模的数据集和检索需求，需要选择合适的索引结构以及设置合理的搜索参数以优化相似度搜索效果。

要改善结果相关性，可以尝试以下改进措施：

• 使用专门为地址设计或调整过的NLP模型。
• 在将地址输入BERT之前，确保对地址进行了标准化处理。
• 使用余弦相似度代替欧氏距离来衡量两个向量的相似性。
• 如果数据量允许，考虑使用更高级的索引结构，例如IVF或者HNSW，并调整其参数以获得更好的召回率和精度。

另外，对于地址匹配问题，有时候专门的地址解析和匹配算法会比通用的NLP模型更为有效。