faiss 向量库应用

requirements: pandas、text2vec、faiss-cpu、numpy
可以安装faiss-gpu,运行速度更快

任务描述：提供了两个文件，第一个文件包括Query，第二个包括参考问题和答案(实际为参考问题对应的知识)。我希望能够让Query和对应的知识匹配到一起，用作后续的prompt生成。

Queries.csv内容示例：

怎么查询PUK
钥匙扣怎么领啊？
直播中了流量怎么领
搬家了拆机哪里办
到期了怎么续约

参考文档内容示例：

问题ID	类别	问题	答案
27981	号卡服务	您好，我的套餐中有一张副卡没开通，请问要怎么重新开通？	实名认证、补录入和二次实名办理指南 -用户点击此链接：https://workwx.sh.189.cn/work/weixin/file/2022-06-27_e7204c9d-c88a-4661-9383-88d5f386f9f1.png 打开图片，扫码后可进行实名认证、补录入和二次实名。 -收到号卡激活: 用户需提供身份证原件，快递员或专员上门面对面激活，或通过……
27980	营销权益活动产品	这个免费30G不是免费的吗？后续还要额外购买收费？	宠粉有礼（简称“宠粉流量”、“0元5G”、“0元10G”、“企微新人礼”、“新人礼”、“新人领10G”、“新人领30G”） -宠粉有礼活动为上海电信为首次添加专属通信管家的手机用户，仅限上海地区用户本网用户提供的活动，本活动仅针对首次添加专属通信管家的用户，同一……

解决思路：Query与参考的问题做相似度匹配，根据匹配到的参考问题，与参考知识做匹配。

要提前准备embedding模型文件(text2vec-bge-large-chinese等)。

生成数据库相关文件.npy、.index。

最终会生成一个匹配文件，包括columns=[‘Query’, ‘Similarity’, ‘问题ID’, ‘问题’, ‘类别’, ‘答案’]

import pandas as pd
from text2vec import SentenceModel
import faiss
import numpy as np

class FAISS:
    def __init__(self, excel_path, embeddings_model_path):
        self.excel_path = excel_path
        self.embedding_model = SentenceModel(embeddings_model_path)
        self.documents = self._load_documents_from_excel()
        self.knowledge_base_embeddings = self._generate_embeddings()
        self.index = self._build_index()

    def _load_documents_from_excel(self):
        # 从Excel文件中加载数据。

        df = pd.read_excel(self.excel_path)
        # 将所有需要的列加载进来

        documents = df[['问题ID', '问题', '类别', '答案']].values.tolist()
        return documents

    def _generate_embeddings(self):
       # 为Excel中的问题生成嵌入向量。

        texts_list = [doc[1] for doc in self.documents]  # 使用问题列生成嵌入

        return self.embedding_model.encode(texts_list)

    def _build_index(self):
        # 使用FAISS构建基于内积的索引，适用于余弦相似度计算。

        d = self.knowledge_base_embeddings.shape[1]  # 向量维度

        index = faiss.IndexFlatIP(d)  # 初始化基于内积的索引

        faiss.normalize_L2(self.knowledge_base_embeddings)
        index.add(self.knowledge_base_embeddings)
        return index

    def save_embeddings_and_index(self, embeddings_path, index_path):
        # 保存嵌入向量和FAISS索引到文件。

        np.save(embeddings_path, self.knowledge_base_embeddings)
        faiss.write_index(self.index, index_path)

    def _save_document_names(self, names_path):
        # 将文档名称保存到文本文件，每个名称占一行。

        with open(names_path, 'w', encoding='utf-8') as f:
            for doc_name in [doc[1] for doc in self.documents]:
                f.write(doc_name + '\n')

    def save_all(self, embeddings_path, index_path, names_path):
        # 保存嵌入向量、FAISS索引和文档名称列表。

        self.save_embeddings_and_index(embeddings_path, index_path)
        self._save_document_names(names_path)

    def retrieve_knowledge(self, query, top_k, threshold):
        # 根据查询检索最相似的文档，并返回完整信息。

        query_vector = self.embedding_model.encode([query])
        if query_vector.ndim == 1:
            query_vector = query_vector[np.newaxis, :]
        faiss.normalize_L2(query_vector)
        distances, index_list = self.index.search(query_vector, top_k)
        results = []
        for idx in range(top_k):
            similarity = distances[0][idx]
            if similarity > threshold:
                # 获取完整的文档信息

                doc_info = self.documents[index_list[0][idx]]
                results.append((query, similarity, doc_info[0], doc_info[1], doc_info[2], doc_info[3]))
        return results

# 以下是主程序部分

if __name__ == "__main__":
    excel_path = '副本参考问题.xlsx'  # 替换为你的Excel文件的实际路径

    embeddings_model_path = "text2vec-bge-large-chinese"  # 词嵌入模型路径

    embeddings_path = "test.npy"  # 嵌入向量文件保存路径

    index_path = "test.index"  # FAISS索引文件保存路径

    names_path = "test_names.txt"  # 文档名称文件保存路径  没有用，可以删掉

    faiss_indexer = FAISS(excel_path, embeddings_model_path)
    faiss_indexer.save_all(embeddings_path, index_path, names_path)  # 保存所有数据

    # 进行批量查询

    queries_path = 'queries.csv'  # 查询语句CSV文件路径

    queries_df = pd.read_csv(queries_path, header=None)
    queries = queries_df[0].tolist()  # 读取所有行作为查询语句

    # 定义不同的参数

    settings = [
        {'sheet_name': 'sheet1', 'top_k': 5, 'threshold': 0.5},
     #   {'sheet_name': 'sheet2', 'top_k': 4, 'threshold': 0.7},

      #  {'sheet_name': 'sheet3', 'top_k': 3, 'threshold': 0.8},

    ]

    # 创建一个ExcelWriter对象

    with pd.ExcelWriter('batch_query_results.xlsx') as writer:
        for setting in settings:
            batch_results = []
            for query in queries:
                retrieved_knowledge = faiss_indexer.retrieve_knowledge(query, setting['top_k'], setting['threshold'])
                for result in retrieved_knowledge:
                    batch_results.append(result)

            # 将批量查询结果保存到不同的sheet中

            results_df = pd.DataFrame(batch_results, columns=['Query', 'Similarity', '问题ID', '问题', '类别', '答案'])
            results_df.to_excel(writer, sheet_name=setting['sheet_name'], index=False)

    print("批量查询完成，结果已保存到 batch_query_results.xlsx")

如果是faiss-cpu，向量库的生成可能需要10几分钟，后续进行检索可以直接导入已经保存的向量库数据。

import pandas as pd
from text2vec import SentenceModel
import faiss
import numpy as np

class FAISS:
    def __init__(self, excel_path, embeddings_model_path, embeddings_path, index_path):
        self.excel_path = excel_path
        self.embedding_model = SentenceModel(embeddings_model_path)
        self.embeddings_path = embeddings_path
        self.index_path = index_path
        self.documents = self._load_documents_from_excel()
        self.knowledge_base_embeddings = self._load_embeddings()
        self.index = self._load_index()

    def _load_documents_from_excel(self):
        '''
        从Excel文件中加载数据。
        '''
        df = pd.read_excel(self.excel_path)
        # 将所有需要的列加载进来

        documents = df[['问题ID', '问题', '类别', '答案']].values.tolist()
        return documents

    def _load_embeddings(self):
        '''
        从文件中加载预先生成的嵌入向量。
        '''
        return np.load(self.embeddings_path)

    def _load_index(self):
        '''
        从文件中加载预先生成的FAISS索引。
        '''
        index = faiss.read_index(self.index_path)
        return index

    def retrieve_knowledge(self, query, top_k, threshold):
        '''
        根据查询检索最相似的文档，并返回完整信息。
        '''
        query_vector = self.embedding_model.encode([query])
        if query_vector.ndim == 1:
            query_vector = query_vector[np.newaxis, :]
        faiss.normalize_L2(query_vector)
        distances, index_list = self.index.search(query_vector, top_k)
        results = []
        for idx in range(top_k):
            similarity = distances[0][idx]
            if similarity > threshold:
                # 获取完整的文档信息

                doc_info = self.documents[index_list[0][idx]]
                results.append((query, similarity, doc_info[0], doc_info[1], doc_info[2], doc_info[3]))
        return results

# 以下是主程序部分

if __name__ == "__main__":
    excel_path = '副本参考问题.xlsx'  # 替换为你的Excel文件的实际路径

    embeddings_model_path = "text2vec-bge-large-chinese"  # 词嵌入模型路径

    embeddings_path = "test.npy"  # 预先生成的嵌入向量文件路径

    index_path = "test.index"  # 预先生成的FAISS索引文件路径

    queries_path = 'queries.csv'  # 查询语句CSV文件路径

    faiss_indexer = FAISS(excel_path, embeddings_model_path, embeddings_path, index_path)

    # 进行批量查询

    queries_df = pd.read_csv(queries_path, header=None)
    queries = queries_df[0].tolist()  # 读取所有行作为查询语句

    batch_results = []
    top_k = 3  # 设置检索数量

    threshold = 0.8  # 设置相似度阈值

    for query in queries:
        retrieved_knowledge = faiss_indexer.retrieve_knowledge(query, top_k, threshold)
        for result in retrieved_knowledge:
            batch_results.append(result)

    # 将批量查询结果保存到XLSX文件

    results_df = pd.DataFrame(batch_results, columns=['Query', 'Similarity', '问题ID', '问题', '类别', '答案'])
    results_df.to_excel('query_results.xlsx', index=False)

    print("查询完成，结果已保存到 query_results.xlsx")

query_results.xlsx示例：

Query	Similarity	问题ID	问题	类别	答案
IPTV遥控器坏了怎么报修	0.807556152	1664	IPTV遥控器突然不工作了，换了电池也没用，该怎么办？	自助报障	IPTV报障 -若用户的IPTV电视产生诸如无法登录主页面、直播频道异常、点播/回看异常…… -若用户仍有报障的问题和疑问，可以在聊天框内输入“服务升级”以转接人工客服咨询详情。……

利用 faiss 将 Query 与参考问题做匹配

CATALOG

FEATURED TAGS