FAISS构建知识库3

出现分数很高的问题主要在于 未对嵌入向量进行适当归一化，导致使用内积度量时分数远超出 0 到 1 的范围。如果分数要与原始工程代码一致（即在 0 到 1 之间），需要对向量在添加和查询时都进行 L2 归一化。

在 FAISS 中，使用 IndexFlatIP 时，将向量进行 L2 归一化，内积度量结果会等同于余弦相似度，范围为 [0, 1]。下面是对代码的优化，使相似度分数在 0 到 1 之间。

优化后的代码

import faiss
import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch


class VectorStore:
    def __init__(self, embedding_model: str, embedding_dim: int = 768, distance_metric: str = "ip"):
        """
        初始化向量存储，选择距离度量方式。
        :param embedding_model: 嵌入模型名称，例如 "m3e-base" 或 "bge-large"
        :param embedding_dim: 嵌入维度
        :param distance_metric: 距离度量方式，支持 "l2" (欧氏距离) 或 "ip" (内积)
        """
        self.tokenizer = AutoTokenizer.from_pretrained(embedding_model)
        self.model = AutoModel.from_pretrained(embedding_model)
        self.embedding_dim = embedding_dim

        # 初始化FAISS索引
        if distance_metric == "ip":
            self.index = faiss.IndexFlatIP(embedding_dim)  # 内积度量
        else:
            self.index = faiss.IndexFlatL2(embedding_dim)  # 默认为欧氏距离
        self.distance_metric = distance_metric

    def embed_text(self, texts: list) -> np.ndarray:
        """
        将文本列表转换为嵌入向量
        :param texts: 文本列表
        :return: 归一化的嵌入向量数组
        """
        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
            embeddings = outputs.last_hidden_state[:, 0, :]  # 使用CLS向量
            embeddings = embeddings.cpu().numpy()  # 转为numpy数组

        # 对向量进行L2归一化
        faiss.normalize_L2(embeddings)
        return embeddings

    def add_documents(self, docs: list):
        """
        将文档列表添加到FAISS索引中
        """
        embeddings = self.embed_text(docs)
        self.index.add(np.array(embeddings, dtype=np.float32))  # 添加归一化后的向量到索引
        self.docs = docs  # 保存文档内容以便打印

    def search(self, query: str, top_k: int = 5, score_threshold: float = 0.5):
        """
        在索引中搜索与查询最匹配的文档，并打印所有文档的相似度得分
        """
        print(f"\n查询内容: {query}")

        query_embedding = self.embed_text([query])  # 将查询转换为归一化的嵌入向量

        # 执行检索，获取所有文档的内积（相似度）和索引
        distances, indices = self.index.search(query_embedding, len(self.docs))

        results = []
        print("\n所有文档的相似度得分：")
        for dist, idx in zip(distances[0], indices[0]):
            score = dist  # 由于内积相似度等同于归一化后的余弦相似度，分数直接为距离
            is_valid = score >= score_threshold  # 余弦相似度范围 [0, 1]，分数越高越相似

            # 打印所有文档的得分
            print(f"文档: {self.docs[idx]}, 相似度得分: {score}")

            if is_valid:
                results.append({'doc_index': idx, 'doc': self.docs[idx], 'score': score})

        print("\n匹配成功的文档（符合 score_threshold 的文档）：")
        for res in results:
            print(f"匹配文档索引: {res['doc_index']}, 内容: {res['doc']}, 匹配得分: {res['score']}")
        return results

    def save_index(self, file_path: str):
        """
        保存FAISS索引到文件
        """
        faiss.write_index(self.index, file_path)

    def load_index(self, file_path: str):
        """
        从文件加载FAISS索引
        """
        self.index = faiss.read_index(file_path)


# 使用示例
if __name__ == "__main__":
    docs = ["你好，这是一个样本文档。", "今天的天气很好。", "这是关于AI的介绍。", "机器学习是人工智能的一个分支。"]
    query = "什么是人工智能？"

    # 使用 "m3e-base" 嵌入模型初始化向量库
    vector_store = VectorStore(embedding_model="m3e-base", distance_metric="ip")

    # 添加文档
    vector_store.add_documents(docs)

    # 查询检索
    results = vector_store.search(query, top_k=3, score_threshold=0.5)

    # 保存索引
    vector_store.save_index("faiss_index.bin")

    # 加载索引
    # vector_store.load_index("faiss_index.bin")

关键改动说明

1. L2 归一化：在 embed_text 方法中，对所有嵌入向量进行 L2 归一化，这样在使用内积（ip）度量时，分数会在 0 到 1 之间，与余弦相似度一致。

2. 分数计算与过滤：

   • 由于使用内积度量并归一化嵌入，distances 直接表示余弦相似度。
   • 设置的 score_threshold 可以直接在 0 到 1 的范围内筛选符合条件的文档。

验证

在此优化后，您可以看到所有文档与查询的相似度分数应该会在 0 到 1 之间，与原始工程的逻辑一致。