TF(词频):一个词在当前文档中出现的频率。越高越重要。
IDF(逆文档频率):一个词在整个文档集合中出现的稀有程度。出现越少(如“人工智能”只在少数文档出现),IDF 越高。
TF-IDF = TF × IDF:既考虑词在当前文档的重要性,又考虑它在全局中的独特性。
import jieba
import requests
import json
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
# ==============================
# 配置 Ollama 本地模型
# ==============================
OLLAMA_URL = "http://localhost:11434/api/generate"
MODEL_NAME = "deepseek-r1:7b" # 你拉取的模型名称
def call_ollama(prompt: str, temperature: float = 0.7) -> str:
"""调用本地 Ollama 的 DeepSeek 模型生成文本"""
payload = {
"model": MODEL_NAME,
"prompt": prompt,
"stream": False,
"options": {
"temperature": temperature
}
}
try:
response = requests.post(OLLAMA_URL, json=payload, timeout=120)
response.raise_for_status()
result = response.json()
return result.get("response", "")
except Exception as e:
print(f"调用 Ollama 出错: {e}")
return ""
# ==============================
# 中文分词函数
# ==============================
def chinese_tokenizer(text: str) -> str:
words = jieba.lcut(text)
return " ".join(words)
# ==============================
# 准备文档集合和查询
# ==============================
documents = [
"人工智能正在改变世界",
"AI正在革新包括医疗和金融在内的各个行业",
"机器学习是人工智能的一个子领域",
"深度学习在图像识别方面取得了巨大成功"
]
query = "人工智能如何改变世界?"
# ==============================
# 预处理:分词
# ==============================
tokenized_docs = [chinese_tokenizer(doc) for doc in documents]
tokenized_query = chinese_tokenizer(query)
# ==============================
# TF-IDF 特征提取与检索
# ==============================
vectorizer = TfidfVectorizer()
doc_tfidf = vectorizer.fit_transform(tokenized_docs)
query_tfidf = vectorizer.transform([tokenized_query])
# 计算相似度,取最相关的文档
similarities = cosine_similarity(query_tfidf, doc_tfidf).flatten()
best_idx = similarities.argmax()
best_doc = documents[best_idx]
best_score = similarities[best_idx]
print(f"查询:{query}")
print(f"最相关的文档(相似度 {best_score:.4f}):{best_doc}\n")
# ==============================
# 构造 Prompt,调用本地 DeepSeek 生成扩展
# ==============================
original_text = query # 也可以单独指定
prompt_template = f"""请基于以下参考信息,扩展用户提供的文本。扩展内容应至少包含150个词,提供不同的视角和细节。
【参考信息】
{best_doc}
【用户文本】
{original_text}
【扩展要求】
请结合参考信息,写出详细的扩展内容。"""
print("正在调用本地 DeepSeek-R1:7B 生成扩展文本...")
extended_text = call_ollama(prompt_template, temperature=0.7)
print("\n扩展后的文本:")
print(extended_text)