DPR(Dense Passage Retrieval)是一种稠密检索方法,核心思想是将文本(查询和文档)映射到同一个稠密向量空间,然后通过计算向量之间的相似度(如余弦相似度)来判断文本的相关性。与传统的BM25等稀疏检索方法相比,DPR能更好地捕捉语义信息。
import requests
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
# Ollama 服务地址(默认端口11434)
OLLAMA_URL = "http://localhost:11434/api/embeddings"
EMBED_MODEL = "nomic-embed-text:latest"
def get_embedding(text):
"""通过 Ollama 获取文本的向量表示"""
response = requests.post(
OLLAMA_URL,
json={"model": EMBED_MODEL, "prompt": text}
)
if response.status_code == 200:
return np.array(response.json()["embedding"]).reshape(1, -1)
else:
raise Exception(f"Ollama error: {response.status_code}")
# 查询和文档(中文)
query = "人工智能如何影响科技?"
documents = [
"猫坐在垫子上。",
"狗对着猫叫。",
"敏捷的棕色狐狸跳过懒惰的狗。",
"人工智能正在改变科技行业。",
"机器学习和深度学习是人工智能的子集。",
"由于经济不确定性,今年股市一直波动。",
"人工智能以多种方式影响科技。",
"机器学习算法正在迅速改进。",
"科技领域中人工智能的未来看起来很有希望。"
]
# 获取查询向量
query_emb = get_embedding(query)
# 计算每个文档的相似度
for i, doc in enumerate(documents):
doc_emb = get_embedding(doc)
sim = cosine_similarity(query_emb, doc_emb)[0][0]
print(f"文档 {i} 相似度: {sim:.4f}")
print(f"文档 {i}: {doc}\n")
# 找出最相关文档
similarities = []
for doc in documents:
doc_emb = get_embedding(doc)
sim = cosine_similarity(query_emb, doc_emb)[0][0]
similarities.append(sim)
best_idx = similarities.index(max(similarities))
print(f"与查询最相关的文档是文档 {best_idx}:")
print(documents[best_idx])