RAG检索_TFIDF

import json
import pdfplumber
import jieba
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import normalize

# 1. 读取数据集
questions = json.load(open("questions.json"))
pdf = pdfplumber.open("汽车知识手册.pdf")
pdf_content = []
for page_idx in range(len(pdf.pages)):
    pdf_content.append({
        'page': 'page_' + str(page_idx + 1),
        'content': pdf.pages[page_idx].extract_text()
    })

# print('pdf_content: ', pdf_content[0])

# 2. 对文本进行分词化
question_words = [' '.join(jieba.lcut(x['question'])) for x in questions]
pdf_content_words = [' '.join(jieba.lcut(x['content'])) for x in pdf_content]

# 3. 提取TFIDF
tfidf = TfidfVectorizer()
tfidf.fit(question_words + pdf_content_words) # 构建词典 + 计算idf过程

# 4. 转换为矩阵
question_feat = tfidf.transform(question_words) 
pdf_content_feat = tfidf.transform(pdf_content_words)

# 5. 归一化，最终相似度打分在规整范围中
question_feat = normalize(question_feat)
pdf_content_feat = normalize(pdf_content_feat)

# 6. 通过TFIDF进行检索

# 6.1 检索top1
for query_idx, feat in enumerate(question_feat):
    # 对每一个提问 与每页PDF进行打分
    score = feat @ pdf_content_feat.T
    score = score.toarray()[0]
    max_score_page_idx = score.argsort()[::-1][0] + 1
    questions[query_idx]['reference'] = 'page_' + str(max_score_page_idx)

output_url_top1 = 'submit_tfidf_retrieval_top1.json' # 保存top1检索结果
with open(output_url_top1, 'w', encoding='utf8') as up:
    json.dump(questions, up, ensure_ascii=False, indent=4)
    print('top1检索完成，结果已保存到', output_url_top1)

# 6.2 检索top11
for query_idx, feat in enumerate(question_feat):
    score = feat @ pdf_content_feat.T
    score = score.toarray()[0]
    max_score_page_idx = score.argsort()[::-1] + 1
    questions[query_idx]['reference'] = ['page_' + str(x) for x in max_score_page_idx[:10]]

output_url_top11 = 'submit_tfidf_retrieval_top11.json'
with open(output_url_top11, 'w', encoding='utf8') as up:
    json.dump(questions, up, ensure_ascii=False, indent=4)
    print('top11检索完成，结果已保存到', output_url_top11)