Python RAG DPR

作者：追风剑情发布于：2026-5-18 16:49 分类：AI

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])

运行测试