bge_embedding/patch_intent_cls_local.py

# 从本地加载模型进行推理，bge和reranker均使用本地模型
from FlagEmbedding import FlagAutoModel, FlagReranker
from intent_description_template import template, intent_code
import torch
import os
import faiss
import numpy as np
import re

script_dir = os.path.dirname(os.path.abspath(__file__))


class IntentRecognizer:

    def __init__(self):
        # 加载元数据
        self.template_meta = {}
        for k, v in template.items():
            for desc in v:
                self.template_meta[desc] = k
        # 模板元数据
        self.template_meta_list = list(self.template_meta.keys())
        # 加载模型
        self.model = FlagAutoModel.from_finetuned(os.path.join(script_dir, "bge-m3"),
                                             query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                                             local_files_only=True,
                                             use_fp16=True,
                                             pooling_method="cls",
                                             devices=["cuda:0"])
        self.reranker = FlagReranker(os.path.join(script_dir, 'bge-reranker-v2-m3'), use_fp16=True, local_files_only=True, devices=["cuda:0"])
        # 模型预热
        print("模型预热中...")
        self.model.encode(["这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。"])
        self.reranker.compute_score([("这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。",
                                 "这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。")])
        self.print_gpu()
        # 加载向量数据库
        self.database_index = None
        database_path = os.path.join(script_dir, "intent_index.faiss")
        if not os.path.exists(database_path):
            # 要求embeddings是一个二维矩阵，类型为float32
            embeddings = self.model.encode(self.template_meta_list)['dense_vecs'].astype(np.float32)  # 选取密集向量，变为float32
            faiss.normalize_L2(embeddings)  # L2归一化
            # Create FAISS index
            dimension = embeddings[0].shape[0]
            self.database_index = faiss.IndexFlatIP(dimension)  # 建立内积索引
            self.database_index.add(embeddings)  # 添加索引
            # Save for future use
            faiss.write_index(self.database_index, database_path)

        if self.database_index is None:
            self.database_index = faiss.read_index(database_path)

    @staticmethod
    def print_gpu():
        if torch.cuda.is_available():
            print(f"allocated:{torch.cuda.memory_allocated()/1024**3:.2f}GB", end=' ')
            print(f"reserved: {torch.cuda.memory_reserved()/1024**3:.2f}GB")

    # 应用推理阶段
    def pick_out(self, query, top_k):
        # 要求query_embedding是一个二维矩阵，形状为(1, 1024)
        query_embedding = self.model.encode([query])['dense_vecs'].astype(np.float32)
        faiss.normalize_L2(query_embedding)
        distances, indices = self.database_index.search(query_embedding, top_k)
        group_query = [(query, self.template_meta_list[indices[0][i]]) for i in range(top_k)]
        score = self.reranker.compute_score(group_query, normalize=True)
        rerank_result = sorted([(distances[0][_], indices[0][_], score[_]) for _ in range(top_k)], key=lambda x: x[2],
                               reverse=True)  # distance, indices, rerank_score
        score_idx = 2  # 重排序相关度
        meta_idx = 1  # 模板位置
        similarity_idx = 0  # 向量相似度
        print("***检索结果***：")
        for i in range(top_k - 1, -1, -1):
            print(
                f"***{i} 相关度:{rerank_result[i][score_idx]:.2f} 相似度:{rerank_result[i][similarity_idx]:.2f} 意图:{self.template_meta.get(self.template_meta_list[rerank_result[i][meta_idx]])} 关联:{self.template_meta_list[rerank_result[i][meta_idx]]}")

        # 从排序结果中拆解到意图的大小类编号
        result = []  # 意图识别结果
        confidence = []# 置信度
        for i in range(top_k):
            # 拿到描述词
            description = self.template_meta_list[rerank_result[i][meta_idx]]
            # 拿到自编码
            custom_number = self.template_meta[description]
            # 拿到大小类标号
            result.append(intent_code[custom_number])
            # 添加置信度
            confidence.append(rerank_result[i][score_idx])

        return confidence, result


def quick_answer_q_a_v2(question, is_english=0):
    """快速问答分支，赋值metric"""
    metric = ''
    question += ' ' # 深拷贝
    if is_english == 0:
        pattern_res_open_qa = re.findall("(开启|打开).*?问答", question)
        pattern_res_close_qa = re.findall("(关闭|关掉).*?问答", question)
    else:
        question = question.strip().lower()
        pattern_res_open_qa = re.findall(
            r'\b(?:open|show me|enable)\b\W*.*?(?:quiz\W+with\W+prizes|enable\W+award\W*-\W*winning\W+q\W*&\W*a)',
            question)
        pattern_res_close_qa = re.findall(
            r'\b(?:close|disable)\b\W*.*?(?:quiz\W+with\W+prizes|enable\W+award\W*-\W*winning\W+q\W*&\W*a)', question)

    if len(pattern_res_open_qa) > 0:
        metric = "openQandA"

    if len(pattern_res_close_qa) > 0:
        metric = "closeQandA"

    return metric

# 单例模式
recognizer_bge = IntentRecognizer()