Python RAG ANN

作者：追风剑情发布于：2026-5-19 14:44 分类：AI

ANN (approximate nearest neighbor，近似最近邻)，FAISS 库支持十亿级数量级文档查询。下面是采用 FAISS 库实现 ANN 的示例。

import faiss
import numpy as np
import requests

# ========== 使用本地 Ollama 替代 BERT ==========
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"])
    else:
        raise Exception(f"Ollama error: {response.status_code}")

def encode(texts):
    """对文本列表进行编码，返回稠密向量矩阵"""
    embeddings = [get_embedding(text) for text in texts]
    return np.array(embeddings)

def normalize_vectors(vectors):
    """归一化向量，使得余弦相似度可以用内积计算"""
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    return vectors / (norms + 1e-10)  # 防止除零错误

# 文档库（中文）
documents = [
    "猫坐在垫子上。",
    "狗对着猫叫。",
    "敏捷的棕色狐狸跳过懒惰的狗。",
    "人工智能正在改变科技行业。",
    "机器学习和深度学习是人工智能的子集。",
    "由于经济不确定性，今年股市一直波动。",
    "人工智能以多种方式影响科技。",
    "机器学习算法正在迅速改进。",
    "科技领域中人工智能的未来看起来很有希望。"
]

# 用户查询（中文）
query = "人工智能如何影响科技？"

# 对文档和查询进行编码
print("正在编码文档...")
document_embeddings = encode(documents)
print("正在编码查询...")
query_embedding = encode([query])[0]

# 归一化向量（使内积等价于余弦相似度）
document_embeddings = normalize_vectors(document_embeddings)
query_embedding = normalize_vectors(query_embedding.reshape(1, -1))[0]

# 创建 FAISS 索引
dimension = document_embeddings.shape[1]  # 向量维度（nomic-embed-text 是 768）
index = faiss.IndexFlatIP(dimension)  # 使用内积度量（归一化后等价于余弦相似度）
index.add(document_embeddings)  # 添加文档向量到索引

# 执行检索
k = 1  # 取最相关的1个文档
distances, indices = index.search(np.array([query_embedding]), k)

# 输出检索结果
print("\n========== FAISS 检索结果 ==========")
for i, idx in enumerate(indices[0]):
    print(f"文档 {idx} 相似度: {distances[0][i]:.4f}")
    print(f"文档 {idx}: {documents[idx]}")

# 可选：检索多个结果
print("\n========== Top-3 检索结果 ==========")
k = 3
distances, indices = index.search(np.array([query_embedding]), k)
for i, idx in enumerate(indices[0]):
    print(f"第{i+1}名 - 文档 {idx} 相似度: {distances[0][i]:.4f}")
    print(f"  文档内容: {documents[idx]}")

运行测试