一种基于 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)