---
license: mit
language:
- en
- zh
base_model:
- Qwen/Qwen2.5-7B
---

# 对数据任务类型分类，比如"情感分析"、"文本分类"、"翻译"，"总结"、"数学问答"....



    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    import json
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from tqdm import tqdm
    from loguru import logger
    
    model_name = "Laurie/Qwen2.5-7b-data-classification"
    
    # 加载模型和 tokenizer，同时调整 padding_side 为 left（适用于 decoder-only 模型）
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype="auto",
        device_map="auto",
        # attn_implementation="flash_attention_2"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") # batch 推理时要左填充
    
    # 对话模板
    system_message = [{"role": "system", "content": "你是一个数据分类专家，请根据对话内容判断其所属的类别。"}]
    last_query = [{"role": "user", "content": "现在请输出你的判断结果："}]
    
    def prepare_text(messages: list[dict]) -> str:
        """
        将 messages 中的 "from"/"value" 键转为 "role"/"content"，并构造完整对话文本
        """
        messages = [{"role": item["from"], "content": item["value"]} for item in messages]
        messages = system_message + messages + last_query
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        return text
    
    def generate_task_types_batch(messages_batch: list[list[dict]]) -> list[str]:
        """
        对一个 batch 的对话列表进行推理生成，并返回每个对话中 assistant 的回答部分
        """
        # 将每个消息列表转换为完整文本
        texts = [prepare_text(messages) for messages in messages_batch]
        
        # 使用批量编码，并进行 padding 以适应批量输入
        model_inputs = tokenizer(
            texts, 
            return_tensors="pt", 
            padding=True, 
            truncation=True
        ).to(model.device)
        
        with torch.no_grad():
            generated_ids = model.generate(
                **model_inputs,
                max_new_tokens=32,
                eos_token_id=[151643, 151645],
                pad_token_id=151643,
                do_sample=True,
                repetition_penalty=1.05, 
                temperature=0.7,
                top_p=0.8,
                top_k=20
            )
        
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]
    
        task_types = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
        return task_types
    
    def process_json(json_path: str, save_path: str, batch_size: int = 8):
        """
        读取 JSON 文件，对数据进行批量推理处理，
        并将结果写回保存。
        """
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        # 分批处理，batch_size 可根据 GPU 显存情况进行调整
        for i in tqdm(range(0, len(data_slice), batch_size)):
            batch = data_slice[i : i + batch_size]
            conversations_batch = [item["conversations"] for item in batch]
            task_types = generate_task_types_batch(conversations_batch)
            for item, answer in zip(batch, task_types):
                item["task_type"] = answer
    
        with open(save_path, "w", encoding="utf-8") as f:
            json.dump(data_slice, f, ensure_ascii=False, indent=4)
        logger.info(f"已处理 {len(data_slice)} 条数据，保存到 {save_path}")
    if __name__ == "__main__":
        json_path = "./qwen_bench_300k.json"
        save_path = "./qwen_bench_300k_cls.json"
        process_json(json_path, save_path, batch_size=16)