鸟语天空
Python RAG Sentence-BERT
post by:追风剑情 2026-5-18 17:21

一种基于 BERT 改进的模型,专门用于生成句子级别的语义向量(Sentence Embedding),使用两个共享权重的 BERT 网络同时处理两个句子,输出向量后计算相似度,相比原始 BERT,Sentence-BERT 可以将相似度计算时间从 65 小时降到约 5 秒(对 1 万个句子)。

import requests
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

OLLAMA_URL = "http://localhost:11434/api/embeddings"
OLLAMA_GENERATE_URL = "http://localhost:11434/api/generate"
EMBED_MODEL = "nomic-embed-text:latest"
GENERATE_MODEL = "qwen2.5:7b"   # 使用明确存在的模型

def get_embedding(text):
    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)
    raise Exception(f"Embedding error: {response.status_code}")

def chat_with_ollama(prompt):
    response = requests.post(
        OLLAMA_GENERATE_URL,
        json={"model": GENERATE_MODEL, "prompt": prompt, "stream": False}
    )
    if response.status_code == 200:
        return response.json()["response"]
    else:
        return f"生成失败 (HTTP {response.status_code}): {response.text}"

# 文档和查询(同上)
documents = [
    "人工智能正在改变世界。",
    "人工智能正在革新包括医疗和金融在内的各个行业。",
    "人工智能对技术产生了深远的影响。",
    "机器学习和深度学习是人工智能的子领域。",
    "人工智能算法正在快速发展。",
    "人工智能在技术领域的未来前景广阔。"
]
query = "人工智能如何影响技术?"

print("正在计算文档向量...")
document_embeddings = np.vstack([get_embedding(doc) for doc in documents])
query_embedding = get_embedding(query)
similarities = cosine_similarity(query_embedding, document_embeddings).flatten()

print("\n========== 语义检索结果 ==========")
for i, score in enumerate(similarities):
    print(f"文档 {i} 相似度: {score:.4f} - {documents[i]}")

most_relevant_index = similarities.argmax()
most_relevant_doc = documents[most_relevant_index]
print(f"\n最相关文档: {most_relevant_doc}")

prompt = f"""请扩展以下文本,提供不同的视角和细节,至少150词。

文本: {most_relevant_doc}

请从技术发展、社会影响、未来趋势等角度进行扩展。"""

print("\n========== 生成扩展内容 ==========")
response = chat_with_ollama(prompt)
print(response)

运行测试
11111.png

评论:
发表评论:
昵称

邮件地址 (选填)

个人主页 (选填)

内容