要在该代码中使用不同的嵌入模型(例如 m3e-base
或 bge-large
),我们可以使代码更具通用性,将嵌入模型作为参数传递,以便能够动态加载任何支持的嵌入模型。此外,我们可以利用 sentence-transformers
或 transformers
库的预训练模型来生产嵌入。
以下是优化后的代码:
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")
代码改进说明
-
嵌入模型动态加载:
- 使用
transformers
库,通过embedding_model
参数传递模型名称,如"m3e-base"
或"bge-large"
等。 AutoTokenizer
和AutoModel
动态加载嵌入模型,使得向量存储能够支持不同的嵌入模型。
- 使用
-
嵌入生成方法:
embed_text
方法将输入文本列表转换为嵌入向量。这里使用last_hidden_state[:, 0, :]
提取每个文本的[CLS]
向量作为句子嵌入。- 对于
l2
度量,执行 L2 归一化,以便后续计算欧氏距离反映出相似度。
-
add_documents
和search
:add_documents
方法生成嵌入并添加到FAISS
索引中,同时保存文档内容,便于在搜索时打印。search
方法支持输出所有文档的相似度分数,并根据score_threshold
筛选出匹配的文档。
-
使用示例:
- 示例展示了如何使用
"m3e-base"
作为嵌入模型初始化向量存储库,并将文档添加到索引。 search
方法展示了与查询最匹配的文档及其相似度得分。- 还提供了将索引保存到文件和从文件加载索引的功能。
- 示例展示了如何使用
优化
为满足您的要求,我们可以将代码优化为同时打印以下内容:
- 查询内容
query
。 - 所有文档与
query
的相似度分数(包括是否匹配成功)。 - 匹配成功的文档(满足
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")
优化细节说明
-
显示所有文档的相似度得分:
- 在
search
方法中,我们对distances
和indices
进行迭代处理,计算出每个文档的相似度分数score
并打印。 - 这包括所有文档的得分,无论是否符合
score_threshold
。
- 在
-
符合
score_threshold
的文档:- 在计算每个文档的
score
后,根据distance_metric
判断相似性度量标准:ip
(内积):直接使用距离作为相似度,分数越高越相似。l2
(欧氏距离):将距离转换为相似度score = 1 / (1 + dist)
,分数越高越相似。
- 对符合
score_threshold
的结果进行过滤,并将这些匹配成功的结果保存到results
列表中,最后打印符合阈值的文档及其相似度得分。
- 在计算每个文档的
-
更好的输出格式:
search
方法的输出格式被分为三部分:查询内容、所有文档的相似度得分、匹配成功的文档。
余弦相似度
上面的代码 并不是通过余弦相似度 来计算 query
和存储数据的相似度。代码中使用的是 欧氏距离 或 内积 作为相似度度量方式:
-
内积(Inner Product):
- 当
distance_metric="ip"
时,使用的是内积度量(点积),分数越高表示越相似。 - 内积并不直接等同于余弦相似度,但在对向量进行单位化(归一化)处理后,内积可以近似为余弦相似度。
- 当
-
欧氏距离(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
的结果。
和工程代码相同
如果您的目标是确保相似度计算与之前工程代码相同,那么我们可以遵循以下方法:
-
确认与之前代码的距离度量方式一致:在您之前提供的代码片段中,
similarity_search_with_score_by_vector
方法中使用了 内积(Inner Product) 和 欧氏距离(L2 Distance) 作为相似度度量方式,而并未指定使用余弦相似度。因此,我们将继续使用这两种度量方式,并通过模型归一化确保计算符合要求。 -
模型实现的归一化处理:根据之前工程代码的逻辑,可以通过对查询和嵌入向量的归一化来确保一致性。对于
distance_strategy
为l2
时,归一化后的内积可以等效于余弦相似度,而对于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
的结果会被过滤并输出。
- 在