心法利器[114] | 通用大模型文本分类实践（含代码）

文本分类在NLP任务里有多重要我就不多说了，之前我也经常提一些比较常用的文本分类方案，比较容易想到的是从fasttext开始，后续的textcnn、bert等系列方案，然后还有以搜代分（心法利器[60] | 以搜代分的生效机理）和词典（心法利器[41] | 我常说的词典匹配到底怎么做）之类的方案吧，大模型出来后，势必要用大模型来试试看。

叠个甲，本文提出了一种相对简便而且baseline还不低的方案，同时开源了代码，供大家尝试使用，不代表药到病除，具体效果需要结合实际情况进行调优（后面我估计会出一期手把手bad case分析的文章，敬请期待）。

目录：

• 基本原理
• 具体实现
• 效果分析
• 改进空间
• 方案机理理解

基本原理

由于大模型自己具备较强的理解和推理能力，常规的指令大模型都是了解的，因此利用大模型做文本分类更关注下面几个内容：

• 分类任务的具体目标需要在prompt中体现。
• 尽可能每个类目的概念都有相对详细的解释，尤其尤其强调类目之间的差别。

而配合in-context learning的思想，比较简洁地使用大模型进行文本分类的prompt应该包含如下成分：

1. 分类任务的介绍及其需求细节。
2. 每个类目的概念解释。
3. 每个类目最好还有些例子（用学术的方法说，就是few-shot吧）。
4. 需要分类的文本。

但在实际应用过程中，可能会出现类目较多、样本较多的问题，2/3是非常容易让prompt膨胀过长的，然而很长的prompt往往容易让大模型的推理效果下降，里面某些内容要求容易被忽略，因此如果有经过筛选再进入大模型就会变得更方便。因此，前面借助向量检索来缩小范围，然后交给大模型来进行最终的决策。

此时方案就形成了，思路如下。

• 离线，提前配置好每个类目的概念及其对应的样本。（某种程度上，其实算是一种训练，整个思路其实就跟KNN里的训练是类似的）
• 在线，先对给定query进行向量召回，然后把召回结果信息交给大模型做最终决策。

这么说比较抽象，这里我给出例子，方便大家理解处理吧。

强调，本方法不对任何模型进行参数更新，都是直接下载开源模型参数直接使用的，这也算是本方案的一大优势吧。

具体实现

代码结构

.
|-- config
|   `-- toutiao_config.py
|-- data
|   |-- index
|   |   `-- vec_index_toutiao_20240629
|   |       |-- forward_index.txt
|   |       `-- invert_index.faiss
|   `-- toutiao_cat_data
|       |-- class_def.tsv
|       |-- test_set_20240629.txt
|       `-- toutiao_cat_data.txt
|-- script
|   |-- build_vec_index.py
|   `-- run_toutiao_cases.py
`-- src
    |-- classifier.py
    |-- models
    |   |-- llm
    |   |   |-- llm_model.py
    |   |   `-- test_qwen.py
    |   `-- vec_model
    |       |-- simcse_model.py
    |       `-- vec_model.py
    |-- searcher
    |   |-- searcher.py
    |   `-- vec_searcher
    |       |-- vec_index.py
    |       `-- vec_searcher.py
    `-- utils
        `-- data_processing.py

解释一下：

• src是核心代码，data内是原始数据和生成的必须数据，config是配置文件，script是必要的批跑脚本。
• 核心代码内，分成了4个部分，classifier是集成好的分类器，models里面存放的是两个模型类，searcher内是检索模块，utils内就是比较普通的工具函数了。

熟悉我的小伙伴应该有发现，整个项目的结构和之前我写的basic_rag（https://github.com/ZBayes/basic_rag）整个项目的非常相近，仔细想想大家也会理解，本文中提及的大模型文本分类方案，其实就是一种RAG，通过检索查询到用户query接近的样本，然后利用大模型来生成最终的类目，这个就是RAG的含义。有关这块的代码，我分了几期来展开讲解了：

• 心法利器[104] | 基础RAG-向量检索模块（含代码）
• 心法利器[105]  基础RAG-大模型和中控模块代码（含代码）
  

有这个理解，看这个项目的整合就会更加清晰了。

models

本模块使用了两个模型，分别是simcse向量表征模型，以及qwen2-1.5B的大模型基座，此处两者都没有进行额外的训练，参数下载后直接使用。

simcse使用的是https://blog.csdn.net/qq_44193969/article/details/126981581提供的加载方案。

import torch
import torch.nn as nn
from loguru import logger
from tqdm import tqdm
from transformers import BertConfig, BertModel, BertTokenizer

class SimcseModel(nn.Module):
    # https://blog.csdn.net/qq_44193969/article/details/126981581
    def __init__(self, pretrained_bert_path, pooling="cls") -> None:
        super(SimcseModel, self).__init__()

        self.pretrained_bert_path = pretrained_bert_path
        self.config = BertConfig.from_pretrained(self.pretrained_bert_path)
        
        self.model = BertModel.from_pretrained(self.pretrained_bert_path, config=self.config)
        self.model.eval()
        
        # self.model = None
        self.pooling = pooling
    
    def forward(self, input_ids, attention_mask, token_type_ids):
        out = self.model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)

        return out.last_hidden_state[:, 0]

这个模型我在外面多包了一层，方便内部进行模型切换，即有一个vec_model。值得注意的是，此处有一个带有v2的方案，这是之前我写的加速方案，此处只有推理部分，完整原文、加速代码以及具体实验可参考：心法利器[107] onnx和tensorRT的bert加速方案记录，这块并非本文重点，就不赘述了。

import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger

from transformers import BertTokenizer

from src.models.vec_model.simcse_model import SimcseModel

import onnxruntime as ort

class VectorizeModel:
    def __init__(self, ptm_model_path, device = "cpu") -> None:
        self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path)
        self.model = SimcseModel(pretrained_bert_path=ptm_model_path, pooling="cls")
        # print(self.model)
        self.model.eval()
        
        self.DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu")
        # self.DEVICE = device
        logger.info(self.DEVICE)
        self.model.to(self.DEVICE)
        
        self.pdist = nn.PairwiseDistance(2)
    
    def predict_vec(self,query):
        q_id = self.tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt')
        with torch.no_grad():
            q_id_input_ids = q_id["input_ids"].squeeze(1).to(self.DEVICE)
            q_id_attention_mask = q_id["attention_mask"].squeeze(1).to(self.DEVICE)
            q_id_token_type_ids = q_id["token_type_ids"].squeeze(1).to(self.DEVICE)
            q_id_pred = self.model(q_id_input_ids, q_id_attention_mask, q_id_token_type_ids)

        return q_id_pred

    def predict_vec_request(self, query):
        q_id_pred = self.predict_vec(query)
        return q_id_pred.cpu().numpy().tolist()
    
    def predict_sim(self, q1, q2):
        q1_v = self.predict_vec(q1)
        q2_v = self.predict_vec(q2)
        sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1)
        return sim.cpu().numpy().tolist()

class VectorizeModel_v2(VectorizeModel):
    def __init__(self, ptm_model_path, onnx_path, providers=['CUDAExecutionProvider']) -> None:
        # ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
        self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path)
        self.model = ort.InferenceSession(onnx_path, providers=providers)
        
        self.pdist = nn.PairwiseDistance(2)
    
    def _to_numpy(self, tensor):
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
    
    def predict_vec(self,query):
        q_id = self.tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt')
        input_feed = {
            self.model.get_inputs()[0].name: self._to_numpy(q_id["input_ids"]),
            self.model.get_inputs()[1].name: self._to_numpy(q_id["attention_mask"]),
            self.model.get_inputs()[2].name: self._to_numpy(q_id["token_type_ids"]),
        }
        return torch.tensor(self.model.run(None, input_feed=input_feed)[0])
    
    def predict_sim(self, q1, q2):
        q1_v = self.predict_vec(q1)
        q2_v = self.predict_vec(q2)
        sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1)
        return sim.numpy().tolist()

if __name__ == "__main__":
    import time,random
    from tqdm import tqdm
    device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
    # device = ""
    # vec_model = VectorizeModel('C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext', device=device)
    vec_model = VectorizeModel_v2('C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext',
                                 "./data/model_simcse_roberta_output_20240211.onnx",providers=['CUDAExecutionProvider'])
    # vec_model = VectorizeModel_v2('C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext',
    #                              "./data/model_simcse_roberta_output_20240211.onnx",providers=['TensorrtExecutionProvider'])
    # 单测
    # q = ["你好啊"]
    # print(vec_model.predict_vec(q))
    # print(vec_model.predict_sim("你好呀","你好啊"))
    tmp_queries = ["你好啊", "今天天气怎么样", "我要暴富"]
    # 开始批跑
    batch_sizes = [1,2,4,8,16]
    for b in batch_sizes:
        for i in tqdm(range(100),desc="warmup"):
            tmp_q = []
            for i in range(b):
                tmp_q.append(random.choice(tmp_queries))
            vec_model.predict_vec(tmp_q)
        for i in tqdm(range(1000),desc="batch_size={}".format(b)):
            tmp_q = []
            for i in range(b):
                tmp_q.append(random.choice(tmp_queries))
            vec_model.predict_vec(tmp_q)

另一方面就是千问模型了，此处我也包装了一层方便使用，里面基本没什么复杂的东西，就是跟着官方教程走，然后划分模块单独弄了一波而已。

# from transformers import AutoModel, AutoTokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Tuple, List
from loguru import logger

class QWen2Model:
    def __init__(self, model_path, config = {}, device="cuda"):
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype="auto",
            device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = self.model.eval()
        self.device = device

        self.generate_config = self._read_config_(config)
        logger.info("load LLM Model done")
    
    def _read_config_(self, config):
        tmp_config = {}
        tmp_config["max_length"] = config.get("max_length", 2048)
        tmp_config["num_beams"] = config.get("num_beams", 1)
        tmp_config["do_sample"] = config.get("do_sample", False)
        tmp_config["top_k"] = config.get("top_k", 1)
        tmp_config["temperature"] = config.get("temperature", 0.8)
        return tmp_config

    def predict(self, query):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": query}
        ]
        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)

        # Directly use generate() and tokenizer.decode() to get the output.
        # Use `max_new_tokens` to control the maximum output length.
        generated_ids = self.model.generate(
            model_inputs.input_ids,
            attention_mask=model_inputs.attention_mask,
            pad_token_id=self.tokenizer.eos_token_id,
            max_new_tokens=512,
            **self.generate_config
        )
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]

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

if __name__ == "__main__":
    from config.toutiao_config import LLM_CONFIG,LLM_PATH
    print(LLM_CONFIG)
    llm_model = QWen2Model(LLM_PATH, config = LLM_CONFIG, device="cuda")
    print(llm_model.predict("如何做番茄炒蛋"))

searcher

检索器和之前的basic_rag类似，低层使用的是FAISS索引工具做向量索引，这里会分3层，分别是index->vec_searcher->searcher，index重在索引的构建，vec_searcher聚焦向量的检索，searcher是综合检索器，理解下来和搜索引擎的3层概念接近，索引-BS（basic search）-AS（advanced search），向量检索支持多种索引构造模式，而向量检索只是整体搜索引擎的一部分而已。现在我开始从内向外展示。

首先是基础的index部分，即基础索引，这里就是直接调的FAISS的接口了，初始化、插入、保存、检索功能都有。

import faiss

class VecIndex:
    def __init__(self) -> None:
        self.index = ""
    
    def build(self, index_dim):
        description = "HNSW64"
        measure = faiss.METRIC_L2
        self.index = faiss.index_factory(index_dim, description, measure)
    
    def insert(self, vec):
        self.index.add(vec)
    
    def batch_insert(self, vecs):
        self.index.add(vecs)
    
    def load(self, read_path):
        # read_path: XXX.index
        self.index = faiss.read_index(read_path)

    def save(self, save_path):
        # save_path: XXX.index
        faiss.write_index(self.index, save_path)
    
    def search(self, vec, num):
        # id, distance
        return self.index.search(vec, num)

然后是向量检索器vec_searcher。另外需要提醒，此处我是把正排放在这一层了，当然直接放到searcher层也是可以的，因为一套正排背后可能有多套索引或者子检索器。接口上，其实和底层的index几乎是一样的，不过对数据的处理会更精细。

import os, json
from loguru import logger
from src.searcher.vec_searcher.vec_index import VecIndex

class VecSearcher:
    def __init__(self):
        self.invert_index = VecIndex() # 检索倒排，使用的是索引是VecIndex
        self.forward_index = [] # 检索正排，实质上只是个list，通过ID获取对应的内容
        self.INDEX_FOLDER_PATH_TEMPLATE = "data/index/{}"

    def build(self, index_dim, index_name):
        self.index_name = index_name
        self.index_folder_path = self.INDEX_FOLDER_PATH_TEMPLATE.format(index_name)
        if not os.path.exists(self.index_folder_path) or not os.path.isdir(self.index_folder_path):
            os.mkdir(self.index_folder_path)

        self.invert_index = VecIndex()
        self.invert_index.build(index_dim)

        self.forward_index = []
    
    def insert(self, vec, doc):
        self.invert_index.insert(vec)
        # self.invert_index.batch_insert(vecs)

        self.forward_index.append(doc)
    
    def save(self):
        with open(self.index_folder_path + "/forward_index.txt", "w", encoding="utf8") as f:
            for data in self.forward_index:
                f.write("{}\n".format(json.dumps(data, ensure_ascii=False)))

        self.invert_index.save(self.index_folder_path + "/invert_index.faiss")
    
    def load(self, index_name):
        self.index_name = index_name
        self.index_folder_path = self.INDEX_FOLDER_PATH_TEMPLATE.format(index_name)

        self.invert_index = VecIndex()
        self.invert_index.load(self.index_folder_path + "/invert_index.faiss")

        self.forward_index = []
        with open(self.index_folder_path + "/forward_index.txt", encoding="utf8") as f:
            for line in f:
                self.forward_index.append(json.loads(line.strip()))
    
    def search(self, vecs, nums = 5):
        search_res = self.invert_index.search(vecs, nums)
        recall_list = []
        for idx in range(nums):
            # recall_list_idx, recall_list_detail, distance
            recall_list.append([search_res[1][0][idx], self.forward_index[search_res[1][0][idx]], search_res[0][0][idx]])
        # recall_list = list(filter(lambda x: x[2] < 100, result))

        return recall_list

最外层就是检索了，这里的除了检索，前期的向量表征也要在这一步完成，再者召回的粗排，我也写在了这一步。

import json,requests,copy
import numpy as np
from loguru import logger
from src.searcher.vec_searcher.vec_searcher import VecSearcher
from src.models.vec_model.vec_model import VectorizeModel

class Searcher:
    def __init__(self, model_path, vec_search_path):
        self.vec_model = VectorizeModel(model_path)
        logger.info("load vec_model done")

        self.vec_searcher = VecSearcher()
        self.vec_searcher.load(vec_search_path)
        logger.info("load vec_searcher done")

    def rank(self, query, recall_result):
        rank_result = []
        for idx in range(len(recall_result)):
            new_sim = self.vec_model.predict_sim(query, recall_result[idx][1][0])
            rank_item = copy.deepcopy(recall_result[idx])
            rank_item.append(new_sim)
            rank_result.append(copy.deepcopy(rank_item))
        rank_result.sort(key=lambda x: x[3], reverse=True)
        return rank_result
    
    def search(self, query, nums=3):
        # logger.info("request: {}".format(query))

        q_vec = self.vec_model.predict_vec(query).cpu().numpy()

        recall_result = self.vec_searcher.search(q_vec, nums)

        rank_result = self.rank(query, recall_result)
        # rank_result = list(filter(lambda x:x[4] > 0.8, rank_result))

        # logger.info("response: {}".format(rank_result))
        return rank_result

if __name__ == "__main__":
    VEC_MODEL_PATH = "C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext"
    VEC_INDEX_DATA = "vec_index_test2023121201"
    searcher = Searcher(VEC_MODEL_PATH, VEC_INDEX_DATA)
    q = "什么人不能吃花生"
    print(searcher.search(q))

分类主函数

然后就是分类的主函数了，先放代码再来解释吧。

import copy
import torch
from loguru import logger

from config.toutiao_config import (VEC_INDEX_DATA, VEC_MODEL_PATH,
                                    LLM_CONFIG, LLM_PATH, PROMPT_TEMPLATE,CLASS_DEF_PATH)
from src.searcher.searcher import Searcher
from src.models.llm.llm_model import QWen2Model
from src.utils.data_processing import load_class_def

class VecLlmClassifier:
    def __init__(self) -> None:
        self.searcher = Searcher(VEC_MODEL_PATH, VEC_INDEX_DATA)
        self.device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
        self.llm = QWen2Model(LLM_PATH, LLM_CONFIG, self.device)
        self.PROMPT_TEMPLATE = PROMPT_TEMPLATE
        self.class_def = load_class_def(CLASS_DEF_PATH)

    def predict(self, query):
        # 1. query预处理
        logger.info("request: {}".format(query))
        # 2. query向量召回
        recall_result = self.searcher.search(query, nums=5)
        # logger.debug(recall_result)

        # 3. 请求大模型
        # 3.1 PROMPT拼接
        request_prompt= copy.deepcopy(self.PROMPT_TEMPLATE)
        # 3.1.1 子模块拼接
        examples = []
        options = []
        options_detail = []
        for item in recall_result:
            tmp_examples = "——".join([item[1][0], item[1][1][5]])
            if tmp_examples not in examples:
                examples.append(tmp_examples)
            opt_detail_str = "：".join(["【" + item[1][1][5] + "】",self.class_def[item[1][1][5]]])
            opt = item[1][1][5]
            if opt not in options:
                options.append(opt)
                options_detail.append(opt_detail_str)
        # options.append("拒识：含义不明或用户query所属类目不在列举内时，分为此类")
        examples_str = "\n".join(examples)
        options_str = "，".join(options)
        options_detail_str = "\n".join(options_detail)

        # 3.1.2 整体组装
        request_prompt = request_prompt.replace("<examples>", examples_str)
        request_prompt = request_prompt.replace("<options>", options_str)
        request_prompt = request_prompt.replace("<options_detail>", options_detail_str)
        request_prompt = request_prompt.replace("<query>", query)
        logger.info(request_prompt)

        # 3.2 请求大模型
        llm_response = self.llm.predict(request_prompt)
        # logger.info("llm response: {}".format(llm_response))

        # 3.3 大模型结果解析
        result = "拒识"
        for option in options:
            if option in llm_response:
                result = option
                break
        # logger.info("parse result: {}".format(result))

        # 4. 返回结果
        return result

if __name__ == "__main__":
    import sys
    vlc = VecLlmClassifier()
    if len(sys.argv) > 1:
        logger.info(vlc.predict("".join(sys.argv[1:])))

提一些关键点：

• 此处需要加载的，是检索器（因为我把向量模型写在检索器里了，所以此处就不需要重复加载，当然写到外面通用化也可以）、Qwen大模型还有一些必要的配置项，注意这里的配置除了检索器、大模型的配置，还有一些prompt相关的配置。
• data_processing里面都是各种数据处理的脚本了，批量、重复的数据处理，包括一些数据加载啥的，我都扔这个文件里了，本文就不赘述了。
• 核心流程我写了完整注释，可以直接看，预处理、向量召回、拼prompt并请求大模型。
• prompt模板写在配置文件里，然后预留好预留位，我这里偷了懒，其实类似<query>之类的东西要写成大写的const方前面的，这些都是预留位，即使placeholder。
• 向量召回内容的解析到prompt组装，这块的活比较琐碎，需要仔细写，避免出错。
• 大模型识别非常简单，但是别忘了后面的解析和校验，避免模型出一些奇怪的结果，要结构化最终再来返回结果。

必要脚本

在script文件夹下有两个脚本，一个是用来灌数据的脚本build_vec_index.py，一个是用来跑测数据结果的批跑脚本run_toutiao_cases.py，我一一展示。

import json,torch,copy,random
from tqdm import tqdm
from loguru import logger
from sklearn.model_selection import train_test_split

from src.utils.data_processing import load_toutiao_data
from src.models.vec_model.vec_model import VectorizeModel
from src.searcher.vec_searcher.vec_searcher import VecSearcher 

if __name__ == "__main__":
    # 0. 必要配置
    VERSION = "20240629"
    VEC_MODEL_PATH = "C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext"
    SOURCE_INDEX_DATA_PATH = "./data/toutiao_cat_data/toutiao_cat_data.txt" # 数据来源：https://github.com/aceimnorstuvwxz/toutiao-text-classfication-dataset
    VEC_INDEX_DATA = "vec_index_toutiao_{}".format(VERSION)
    TESE_DATA_PATH = "./data/toutiao_cat_data/test_set_{}.txt".format(VERSION)
    RANDOM_SEED = 100
    # MODE = "DEBUG"
    MODE = "PRO"

    DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu")
    TEST_SIZE = 0.1
    # 类目体系
    CLASS_INFO = [
        ["100", '民生-故事', 'news_story'],
        ["101", '文化-文化', 'news_culture'],
        ["102", '娱乐-娱乐', 'news_entertainment'],
        ["103", '体育-体育', 'news_sports'],
        ["104", '财经-财经', 'news_finance'],
        # ["105", '时政 新时代', 'nineteenth'],
        ["106", '房产-房产', 'news_house'],
        ["107", '汽车-汽车', 'news_car'],
        ["108", '教育-教育', 'news_edu' ],
        ["109", '科技-科技', 'news_tech'],
        ["110", '军事-军事', 'news_military'],
        # ["111" 宗教 无，凤凰佛教等来源],
        ["112", '旅游-旅游', 'news_travel'],
        ["113", '国际-国际', 'news_world'],
        ["114", '证券-股票', 'stock'],
        ["115", '农业-三农', 'news_agriculture'],
        ["116", '电竞-游戏', 'news_game']
    ]
    ID2CN_MAPPING = {}
    for idx in range(len(CLASS_INFO)):
        ID2CN_MAPPING[CLASS_INFO[idx][0]] = CLASS_INFO[idx][1]

    # 1. 加载数据、模型
    # 1.1 加载模型
    vec_model = VectorizeModel(VEC_MODEL_PATH, DEVICE)
    index_dim = len(vec_model.predict_vec("你好啊")[0])
    # 1.2 加载数据
    source_index_data = load_toutiao_data(SOURCE_INDEX_DATA_PATH)
    logger.info("load data done: {}".format(len(source_index_data)))
    if MODE == "DEBUG":
        random.shuffle(source_index_data)
        source_index_data = source_index_data[:100000]
    source_index_data_new = []
    for item in source_index_data:
        item[1].append(ID2CN_MAPPING[item[1][1]])
    # 1.3 训练集测试集划分
    train_list, test_list = train_test_split(source_index_data, test_size=TEST_SIZE, random_state=66)

    # 2. 创建索引并灌入数据
    # 2.1 构造索引
    vec_searcher = VecSearcher()
    vec_searcher.build(index_dim, VEC_INDEX_DATA)

    # 2.2 推理向量
    vectorize_result = []
    for q in tqdm(train_list, desc="VEC MODEL RUNNING"):
        vec = vec_model.predict_vec(q[0]).cpu().numpy()
        tmp_result = copy.deepcopy(q)
        tmp_result.append(vec)
        vectorize_result.append(copy.deepcopy(tmp_result))

    # 2.3 开始存入
    for idx in tqdm(range(len(vectorize_result)), desc="INSERT INTO INDEX"):
        vec_searcher.insert(vectorize_result[idx][2], vectorize_result[idx][:2])

    # 3. 保存
    # 3.1 索引保存
    vec_searcher.save()
    # 3.2 测试集保存
    with open(TESE_DATA_PATH, "w", encoding="utf8") as f:
        for item in test_list:
            f.write("_!_".join(item[1]) + "\n")

注释同样写的比较明白了，说白了就是海量数据的预处理后，逐步把数据存入库中，当然这里没忘记把数据按照一定比例分为训练集和测试集（从KNN的角度，入库过程本质就算是一个训练过程了）。大家也可以根据实际业务场景，调整入库的数据策略，例如限定个数等，这个就大家自己写吧。

另一个脚本是批跑脚本，这个并不难看懂，就是一个读数据、预测、计算指标的流程罢了。

from tqdm import tqdm
from sklearn.metrics import classification_report, confusion_matrix
from loguru import logger

from src.classifier import VecLlmClassifier
from src.utils.data_processing import load_toutiao_data

TEST_DATA_PATH = "data/toutiao_cat_data/test_set_20240629.txt"
test_data = load_toutiao_data(TEST_DATA_PATH)

vlc = VecLlmClassifier()
test_list = []
pred_list = []
labels = set()
for i in tqdm(range(len(test_data)), desc="RUNNING TEST"):
    test_list.append(test_data[i][1][5])
    labels.add(test_data[i][1][5])
    pred_list.append(vlc.predict(test_data[i][0]))
labels = list(labels)

logger.info("\n{}".format(classification_report(test_list, pred_list, labels = labels)))
logger.info("\n{}".format(confusion_matrix(test_list, pred_list, labels=labels)))

数据细节和prompt

此处我用的demo数据是头条的新闻标题分类数据，在这里：https://github.com/aceimnorstuvwxz/toutiao-text-classfication-dataset，很多数据处理的工作基本是从这里来的，数据可以说是量大管饱吧。但其实这个类目体系多少还是有些问题，我这里提一下：

• 类目是多分类任务，理论上要求类目之间互斥，但是类目体系下重合的还是不少的，如“国际-国际”和别的新闻很容易混淆，还有“财经-财经”和“证券-股票”。
• 类目还是比较不均匀的，类似“证券-股票”的数据是非常少的，样本少导致前面提到的类目互斥问题更严重。

prompt这块，我给出我设计的prompt。

你是一个优秀的句子分类师，能把给定的用户query划分到正确的类目中。现在请你根据给定信息和要求，为给定用户query，从备选类目中选择最合适的类目。

下面是“参考案例”即被标注的正确结果，可供参考：
<examples>

备选类目：
<options>

类目概念：
<options_detail>

用户query：
<query>

请注意：
1. 用户query所选类目，仅能在【备选类目】中进行选择，用户query仅属于一个类目。
2. “参考案例”中的内容可供推理分析，可以仿照案例来分析用户query的所选类目。
3. 请仔细比对【备选类目】的概念和用户query的差异。
4. 如果用户quer也不属于【备选类目】中给定的类目，或者比较模糊，请选择“拒识”。
5. 请在“所选类目：”后回复结果，不需要说明理由。

所选类目：

解析：

• 开篇给出角色和具体任务。
• 参考案例对应前面提到的in-context learning，之前有过实验表明，随便给例子和给出和用户query相似的例子相比，后者效果更好，这也是前面要用向量召回的关键原因。
• 备选类目来自向量召回的结果，结果里有什么类目我们就把哪些放在备选类目里，这样有效缩小类目范围，简化分类问题，这也是前面要先做检索的重要原因，类目多了prompt很长而且大模型识别也没那么准。
• 类目概念能让大模型更好地理解到类目的概念，对分类肯定是要有收益的。
• 请注意下是任务相关的约束，例如类目约束、拒识兜底等，是结合大模型的输出结果调校得到的，毕竟要约束类目的规范性、类目个数等信息。

效果分析

很自然的就是要看看效果咋样，我这里用的是头条开源的新闻标题分类数据：https://github.com/aceimnorstuvwxz/toutiao-text-classfication-dataset。

可以看到，在数据比较完善的情况下，效果还是挺高的，但是随着样本的下降，效果衰减的还是很明显，few-shot的能力体现的并不优秀。

有关这里，做一下预告，后续我会有专门的文章将这个的case分析，用来作为case分析的案例。

改进空间

有关这个数据、问题的改进，我感觉可以以此为例专门讲一下bad case分析怎么做，而且很特异化，我这里就不展开了，后续专门写文章聊。

这里，我专门讲讲这套方案本身可以考虑的空间。

• 首先是最容易想到的，就是模型更新，向量模型方面，simcse显然不是现在的sota，类似BGE等的方案，都是可以考虑的；大模型这里，是受限于我自己的电脑问题所以用的比较小，更好更大的大模型还是可以做一下实验的。
• 向量模型会存在两个特殊的极端情况：全都是一个类目，以及召回多个类都没有正确的。对于前者，如果比较信任向量模型，则可以考虑不过大模型了，降低成本；对于后者，需要集中精力优化向量模型，必要时可能就要优化了。
• 如果向量模型信任感不足，且数据比较足，可以考虑向量模型后再接一个更可靠的交互式相似度模型做精排，然后再进大模型，此时向量召回的数据量可以一定程度提升，精排再来压缩到合适的程度。
• 入库样本的典型性和类目概念的清晰度、完整性，都对分类效果有很大影响，精雕细刻prompt有很大价值。
• 如果因为部署成本、耗时等因素，大模型无法上线，完全考虑，这个方案蒸馏一个小模型来做这个事，例如T5。
• 样本的覆盖率和准确性，对一个系统而言，如果一个案例没见过，那系统大概率就不认识，大模型虽然有较强的泛化能力，但不代表全知全能，尤其是在分类这种边界要求明确，内部信息丰富而又复杂的问题下。

方案机理理解

这块方案探索后再深究，我发现一些比较有意思的理解方式，我在这里逐个解释一下。

• 概念解释和样本案例，对大模型而言实质都在做一件事——给大模型解释类目的边界概念，早年分类模型需要样本进行训练，本质也是这个目标，只是因为目前大模型的理解和推理能力变强，且有很强的生成能力，因此这个事可以通过prompt来解决。
• 向量召回在此处的作用，本质是提前筛选更合适的样本和可能性更高的类目，协助大模型更好地理解边界，且从普遍理性而言，更贴切的例子更有利于理解具体概念和含义。
• 向量召回换个角度，这里就是我之前说的“以搜代分”的操作（心法利器[60] | 以搜代分的生效机理），通过比对相似样本达成分类目标，也可以理解为KNN分类。
• 续上，向量召回如果是以搜代分，大模型则可以看做向量召回后的一种精排，一种精筛，前面的向量召回就是召回（缩小范围）+粗排。此时，一整个分类方案就变成一个相对完整地搜索系统，也呼应了我之前有讲过的把一个任务方案当做系统的想法（心法利器[29] | 把文本分类任务做成一个系统，把分类当做一个系统）。
• 再换个角度，极端的，搜索系统可以看做是一个N分类的问题，这个N等于整个库里的物料条数，只不过类目太多且变更太快，所以才需要做各种召回缩小范围，精排精筛等操作。
• 另外，向量召回本身召回就比较粗，像我之前在搜索系统的时候说过（前沿重器[49] | 聊聊搜索系统2：常见架构），粗排是判断“像不像”的问题，但是“谁更像”还得再要一个精排，进一步提升准确率，所以两者互相帮助，逻辑自洽没什么毛病。
• 如果把整套方案当做是搜索，那大模型的生成，这套分类方案就可以看做一种RAG。

方案优缺点

要说一个方案，就要讨论一下这个方案的优缺点，从而方便我们在后续的任务中进行选择。

首先说一下优点。

• 无训练的高基线。整个项目做下来，不需要进行模型的训练，只需要模型推理的资源，就能到达一个比较高的下限。
• 少样本的高基线。标注经常是业务场景的一个痛点，现在只需要给些例子和类目的解释，就能快速解决问题，便捷度还是很高的，这个方案会比以搜代分方案要再高一点。
• 灵活性高。对经常要做类目个数、边界、样本的变更，会更灵活，该类目配置和增删样本就可以解决了，不像分类模型那么死板要重训。

缺点：

• 上限不会太高，还是不如微调向量模型和大模型。
• 只能是通用领域的知识，对专业领域还是避免不了的领域知识问题。
• 对样本数和覆盖率有一定要求。要想效果好，依旧需要更多样本，可以这么理解，从信息传播角度（心法利器[45] | 模型需要的信息提供够了吗），对于没见过的东西，系统不认识是无法做事的，尤其是类似音乐、文学作品之类的信息分类，这套方案的效果甚至不如花点时间总结总结词典然后用词典匹配。
• 老生常谈的大模型成本和耗时问题。不过如果仅仅是这点，还是可以通过本方案预标注样本后训一个小模型来解决。