FAISS创建知识库2

要在该代码中使用不同的嵌入模型(例如 m3e-basebge-large),我们可以使代码更具通用性,将嵌入模型作为参数传递,以便能够动态加载任何支持的嵌入模型。此外,我们可以利用 sentence-transformerstransformers 库的预训练模型来生产嵌入。

以下是优化后的代码:

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 = "l2"):
        """
        初始化向量存储,选择距离度量方式。
        :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数组

        if self.distance_metric == "l2":
            faiss.normalize_L2(embeddings)  # 如果是 L2,进行归一化

        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]):
            if self.distance_metric == "ip":
                score = dist  # 对于内积,分数直接是距离
                is_valid = score >= score_threshold  # 内积的分数是相似度,分数越高越相似
            else:
                score = 1 / (1 + dist)  # 欧氏距离转换为相似度分数
                is_valid = score >= score_threshold  # 欧氏距离的分数越低越相似
            
            # 打印所有文档的得分
            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. 嵌入模型动态加载

    • 使用 transformers 库,通过 embedding_model 参数传递模型名称,如 "m3e-base""bge-large" 等。
    • AutoTokenizerAutoModel 动态加载嵌入模型,使得向量存储能够支持不同的嵌入模型。
  2. 嵌入生成方法

    • embed_text 方法将输入文本列表转换为嵌入向量。这里使用 last_hidden_state[:, 0, :] 提取每个文本的 [CLS] 向量作为句子嵌入。
    • 对于 l2 度量,执行 L2 归一化,以便后续计算欧氏距离反映出相似度。
  3. add_documentssearch

    • add_documents 方法生成嵌入并添加到 FAISS 索引中,同时保存文档内容,便于在搜索时打印。
    • search 方法支持输出所有文档的相似度分数,并根据 score_threshold 筛选出匹配的文档。
  4. 使用示例

    • 示例展示了如何使用 "m3e-base" 作为嵌入模型初始化向量存储库,并将文档添加到索引。
    • search 方法展示了与查询最匹配的文档及其相似度得分。
    • 还提供了将索引保存到文件和从文件加载索引的功能。

优化

为满足您的要求,我们可以将代码优化为同时打印以下内容:

  1. 查询内容 query
  2. 所有文档与 query 的相似度分数(包括是否匹配成功)。
  3. 匹配成功的文档(满足 score_threshold)及其详细信息。

我们可以进一步简化 search 方法,使其能够一次性打印所有相似度信息并将匹配成功的文档筛选出来。以下是优化后的代码:

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 = "l2"):
        """
        初始化向量存储,选择距离度量方式。
        :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数组

        if self.distance_metric == "l2":
            faiss.normalize_L2(embeddings)  # 如果是 L2,进行归一化

        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]):
            if self.distance_metric == "ip":
                score = dist  # 对于内积,分数直接是距离
                is_valid = score >= score_threshold  # 内积的分数是相似度,分数越高越相似
            else:
                score = 1 / (1 + dist)  # 欧氏距离转换为相似度分数
                is_valid = score >= score_threshold  # 欧氏距离的分数越低越相似
            
            # 打印所有文档的得分
            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. 显示所有文档的相似度得分

    • search 方法中,我们对 distancesindices 进行迭代处理,计算出每个文档的相似度分数 score 并打印。
    • 这包括所有文档的得分,无论是否符合 score_threshold
  2. 符合 score_threshold 的文档

    • 在计算每个文档的 score 后,根据 distance_metric 判断相似性度量标准:
      • ip(内积):直接使用距离作为相似度,分数越高越相似。
      • l2(欧氏距离):将距离转换为相似度 score = 1 / (1 + dist),分数越高越相似。
    • 对符合 score_threshold 的结果进行过滤,并将这些匹配成功的结果保存到 results 列表中,最后打印符合阈值的文档及其相似度得分。
  3. 更好的输出格式

    • search 方法的输出格式被分为三部分:查询内容、所有文档的相似度得分、匹配成功的文档。

余弦相似度

上面的代码 并不是通过余弦相似度 来计算 query 和存储数据的相似度。代码中使用的是 欧氏距离内积 作为相似度度量方式:

  1. 内积(Inner Product)

    • distance_metric="ip" 时,使用的是内积度量(点积),分数越高表示越相似。
    • 内积并不直接等同于余弦相似度,但在对向量进行单位化(归一化)处理后,内积可以近似为余弦相似度。
  2. 欧氏距离(L2 Distance)

    • distance_metric="l2" 时,使用的是欧氏距离度量,分数越低表示越相似。
    • 代码中将欧氏距离转换为相似度分数:score = 1 / (1 + dist),但这也不是余弦相似度。

使用余弦相似度的优化方法

为了真正使用余弦相似度来计算 query 和存储数据的相似度,我们可以对向量进行 L2 归一化,然后使用内积度量。以下是基于余弦相似度的修改方法:

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):
        """
        初始化向量存储,并使用余弦相似度来计算相似度
        :param embedding_model: 嵌入模型名称,例如 "m3e-base" 或 "bge-large"
        :param embedding_dim: 嵌入维度
        """
        self.tokenizer = AutoTokenizer.from_pretrained(embedding_model)
        self.model = AutoModel.from_pretrained(embedding_model)
        self.embedding_dim = embedding_dim
        
        # 使用内积度量,并对嵌入进行归一化来模拟余弦相似度
        self.index = faiss.IndexFlatIP(embedding_dim)  # 内积度量

    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
            
            # 打印所有文档的得分
            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")

    # 添加文档
    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")

优化后的代码说明

  • 余弦相似度实现:通过 faiss.normalize_L2 对所有嵌入向量进行 L2 归一化,这样在计算内积时,结果就等同于余弦相似度。
  • 简化的相似度计算:使用内积直接作为余弦相似度的分数,打印和过滤过程都基于该分数。
  • 打印所有分数:在 search 方法中,打印所有文档的余弦相似度分数,同时筛选符合 score_threshold 的结果。

和工程代码相同

如果您的目标是确保相似度计算与之前工程代码相同,那么我们可以遵循以下方法:

  1. 确认与之前代码的距离度量方式一致:在您之前提供的代码片段中,similarity_search_with_score_by_vector 方法中使用了 内积(Inner Product)欧氏距离(L2 Distance) 作为相似度度量方式,而并未指定使用余弦相似度。因此,我们将继续使用这两种度量方式,并通过模型归一化确保计算符合要求。

  2. 模型实现的归一化处理:根据之前工程代码的逻辑,可以通过对查询和嵌入向量的归一化来确保一致性。对于 distance_strategyl2 时,归一化后的内积可以等效于余弦相似度,而对于 ip 时直接使用内积。

下面是确保与之前工程代码一致的改进代码:

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 = "l2"):
        """
        初始化向量存储,选择距离度量方式。
        :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归一化
        if self.distance_metric == "l2":
            faiss.normalize_L2(embeddings)  # 如果是 L2,进行归一化
        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]):
            if self.distance_metric == "ip":
                score = dist  # 对于内积,分数直接是距离
                is_valid = score >= score_threshold  # 内积的分数是相似度,分数越高越相似
            else:
                score = 1 / (1 + dist)  # 欧氏距离转换为相似度分数
                is_valid = score >= score_threshold  # 欧氏距离的分数越低越相似
            
            # 打印所有文档的得分
            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")

代码优化说明

  • 相似度度量方式的一致性distance_metric 参数允许在 内积欧氏距离 间切换,以与之前代码中的相似度度量方式保持一致。
  • L2 归一化
    • 如果使用欧氏距离(l2),在 embed_text 方法中对嵌入向量执行 L2 归一化,这与之前代码中将 l2 转换为相似度的处理方式一致。
    • 对于内积(ip),则无需归一化,直接使用距离作为相似度分数。
  • 打印和过滤
    • search 方法中,所有文档的相似度分数都将打印,并且符合 score_threshold 的结果会被过滤并输出。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值