Python RAG TF-IDF

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

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)

运行测试