RAG检索_BERT(bge)

 

import json
import pdfplumber
import jieba
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import normalize
from sentence_transformers import SentenceTransformer

# 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() or ""
    })

# 2. BGE （sbert 微调之后的模型， 通过modelscope下载）
model = SentenceTransformer('../models/BAAI/bge-small-zh-v1.5/')
question_sentences = [x['question'] for x in questions]
pdf_content_sentences = [x['content'] for x in pdf_content]

# 方法1: 借助BGE-BERT进行文本编码
question_embeddings = model.encode(question_sentences, normalize_embeddings=True)
pdf_embeddings = model.encode(pdf_content_sentences, normalize_embeddings=True)

# 返回Top1
for query_idx, feat in enumerate(question_embeddings):
    score = np.dot(pdf_embeddings, feat)
    max_score_page_idx = score.argsort()[::-1][0] + 1
    questions[query_idx]['reference'] = 'page_' + str(max_score_page_idx)

with open('submit_bge_retrieval_top1.json', 'w', encoding='utf8') as up:
    json.dump(questions, up, ensure_ascii=False, indent=4)

# 返回Top11
for query_idx, feat in enumerate(question_embeddings):
    score = np.dot(pdf_embeddings, feat)
    max_score_page_idx = score.argsort()[::-1] + 1
    questions[query_idx]['reference'] = ['page_' + str(x) for x in max_score_page_idx[:10]]

with open('submit_bge_retrieval_top11.json', 'w', encoding='utf8') as up:
    json.dump(questions, up, ensure_ascii=False, indent=4)



# 方法2: 借助Jina进行文本编码

# 3.1 准备问题embedding, 答案embedding

# 下载在线模型命令行
# modelscope download --model jinaai/jinaai-embedding-v2-base-zh --local_dir jinaai/jinaai-embeddings-v2-base-zh
 
# 模型及数据准备
# jina-embeddings-v2 使用自定义建模代码，需 trust_remote_code=True 才能正确加载权重
model = SentenceTransformer(
    '../models/jinaai/jina-embeddings-v2-base-zh/',
    trust_remote_code=True,
)
question_sentences = [x['question'] for x in questions]
pdf_content_sentences = [x['content'] for x in pdf_content]

question_embeddings = model.encode(question_sentences, normalize_embeddings=True)
pdf_embeddings = model.encode(pdf_content_sentences, normalize_embeddings=True)

# 返回Top1
for query_idx, feat in enumerate(question_embeddings):
    score = np.dot(pdf_embeddings, feat)
    max_score_page_idx = score.argsort()[::-1][0] + 1
    questions[query_idx]['reference'] = 'page_' + str(max_score_page_idx)

with open('submit_jina_retrieval_top1.json', 'w', encoding='utf8') as up:
    json.dump(questions, up, ensure_ascii=False, indent=4)


# 返回Top11
for query_idx, feat in enumerate(question_embeddings):
    score = np.dot(pdf_embeddings, feat)
    max_score_page_idx = score.argsort()[::-1] + 1
    questions[query_idx]['reference'] = ['page_' + str(x) for x in max_score_page_idx[:10]]

with open('submit_jina_retrieval_top11.json', 'w', encoding='utf8') as up:
    json.dump(questions, up, ensure_ascii=False, indent=4)