【人工智能】QWEN + RAG系统

结合qwen2.5-7b + bge-reranker-base + stella-base-zh-v2 组合成的rag系统

其中测试数据和模型需要自行下载，modelscope就可以。

测试数据链接：Streamlit

# -*- coding: utf-8 -*-

from flask import Flask, request, jsonify
import jieba
import json
import pdfplumber
from rank_bm25 import BM25Okapi
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer
import numpy as np


app = Flask(__name__)

# 加载模型和初始化
device = "cuda" if torch.cuda.is_available() else "cpu"

rerank_tokenizer = AutoTokenizer.from_pretrained('/home/sky/model_data/bge-reranker-base')
rerank_model = AutoModelForSequenceClassification.from_pretrained('/home/sky/model_data/bge-reranker-base').to(device)

model_path = '/home/sky/model_data/Qwen/Qwen2.5-7B-Instruct'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_path, torch_dtype="auto",
                                             device_map="auto").to(device)

sent_model = SentenceTransformer('/home/sky/model_data/stella-base-zh-v2')

pdf_content = []


def load_pdf(knowledge_data_path):
    global pdf_content
    pdf = pdfplumber.open(knowledge_data_path)
    pdf_content = []
    for page_idx in range(len(pdf.pages)):
        text = pdf.pages[page_idx].extract_text()
        new_text = split_text_fixed_size(text, chunk_size=100, overlap_size=5)
        for chunk_text in new_text:
            pdf_content.append({
                'page': 'page_' + str(page_idx + 1),
                'content': chunk_text
            })

    # 构建索引
    global pdf_content_words, bm25, pdf_embeddings
    pdf_content_words = [jieba.lcut(x['content']) for x in pdf_content]
    bm25 = BM25Okapi(pdf_content_words)
    pdf_content_sentences = [x['content'] for x in pdf_content]
    pdf_embeddings = sent_model.encode(pdf_content_sentences, normalize_embeddings=True)


def split_text_fixed_size(text, chunk_size, overlap_size):
    new_text = []
    for i in range(0, len(text), chunk_size - overlap_size):
        end = min(i + chunk_size, len(text))
        start = max(0, i - overlap_size)
        new_text.append(text[start:end])
    return new_text


def get_rank_index(max_score_page_idxs_, question_, pdf_content_):
    pairs = [[question_, pdf_content_[idx]['content']] for idx in max_score_page_idxs_]

    inputs = rerank_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512).to(device)
    with torch.no_grad():
        logits = rerank_model(**inputs, return_dict=True).logits.view(-1).float().cpu().numpy()
        # Apply sigmoid to convert logits to probabilities
        scores = 1 / (1 + np.exp(-logits))  # Sigmoid function

    max_score_idx = scores.argmax()
    max_score = scores[max_score_idx]
    index = max_score_page_idxs_[max_score_idx]
    return max_score, index

def qwen_preprocess(tokenizer_, ziliao, question):
    if ziliao:
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": f"帮我结合给定的资料，回答问题。如果问题答案无法从资料中获得，"
                                        f"输出结合给定的资料，无法回答问题. 如果找到答案, 就输出找到的答案, 资料：{ziliao}, 问题:{question}"},
        ]
    else:
        messages = [
            {"role": "system", "content": "你是由智子引擎研发的AI小助手,可以帮助解答各种领域的问题."},
            {"role": "user", "content": f"问题:{question}"},
        ]


    text = tokenizer_.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs_ = tokenizer_([text], return_tensors="pt").to(device)

    return model_inputs_


@app.route('/qa', methods=['POST'])
def qa():
    data = request.json
    question = data.get('question', '')

    if not question:
        return jsonify({'error': 'No question provided'}), 400

    # 首先进行BM25检索
    doc_scores = bm25.get_scores(jieba.lcut(question))
    bm25_score_page_idxs = doc_scores.argsort()[-10:][::-1]

    # 再进行语义检索
    question_embedding = sent_model.encode([question], normalize_embeddings=True)
    score = question_embedding @ pdf_embeddings.T
    ste_score_page_idxs = score.argsort()[0][-10:][::-1]

    bm25_score, bm25_index = get_rank_index(bm25_score_page_idxs, question, pdf_content)
    ste_score, ste_index = get_rank_index(ste_score_page_idxs, question, pdf_content)

    max_score_page_idx = ste_index if ste_score >= bm25_score else bm25_index
    # 得分
    final_score = ste_score if ste_score >= bm25_score else bm25_score
    # 设定一个阈值来判断是否找到了有效的答案
    threshold = 0.5
    if final_score < threshold:
        # 如果得分低于阈值，直接用Qwen模型回答
        model_inputs = qwen_preprocess(tokenizer, "", question)  # 不提供参考资料
        generated_ids = model.generate(
            model_inputs.input_ids,
            max_new_tokens=128,
            attention_mask=model_inputs.attention_mask,
            pad_token_id=tokenizer.eos_token_id
        )
        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
        return jsonify({
            'answer': response,
            'score': float(final_score),
            'reference_page': "N/A"  # 没有参考页面
        })


    model_inputs = qwen_preprocess(tokenizer, pdf_content[max_score_page_idx]['content'], question)

    generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=128,
        attention_mask=model_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return jsonify({'answer': response,"score": float(final_score),'reference_page': pdf_content[max_score_page_idx]['page']})


if __name__ == '__main__':
    knowledge_data_path = './初赛训练数据集.pdf'
    load_pdf(knowledge_data_path)
    app.run(host='0.0.0.0',port=5201,debug=False)