ykallan/SkuInfo-Qwen2.5-3B-R1

商品信息抽取模型

基于Qwen2.5-3B-Instruct模型，可以使用transformers库直接进行推理，经过Lora训练得到，可以在商品名称中抽取出品牌，型号，主商品等信息，以JSON格式输出。

经过强化学习，加强了JSON格式输出校验，以及品牌、型号、主商品这三个字段的抽取准确率。

生成模型可能会出现幻觉现象，不存在的字段可能也会输出结果，比如型号不存在，可能也会把主商品当做型号。

transformers版本低于4.37.0会报错KeyError: qwen2

# encoding: utf-8

import time
from transformers import AutoModelForCausalLM, AutoTokenizer

import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

pretrained_model = "ykallan/SkuInfo-Qwen2.5-3B-R1"

# model = AutoModelForCausalLM.from_pretrained(pretrained_model, low_cpu_mem_usage=True)
model = AutoModelForCausalLM.from_pretrained(pretrained_model)

tokenizer = AutoTokenizer.from_pretrained(pretrained_model)

model.to(device)


def ner(_name: str) -> str:
    messages = [
        {"role": "system", "content": "在以下商品名称中抽取出品牌、型号、主商品，并以JSON格式返回。"},
        {"role": "user", "content": _name}
    ]

    input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    model_inputs = tokenizer([input_ids], return_tensors="pt", padding=True).to(device)

    generate_config = {
        "max_new_tokens": 128
    }

    generated_ids = model.generate(model_inputs.input_ids, **generate_config)
    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]
    # print(response)
    return response


sku_name_list = ["齐心(Comix) 10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 60mm 可装600张 老款",
                 "10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 60mm 可装600张",
                 "齐心10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 30mm 可装300张",
                 "SPC 韩国SPC移动平台打孔机，FP（filepecker）-I（X） 售卖规格：1台",
                 "国保 分体指纹保密柜，Z168-V4 符合国家保密标准(BMB54-2020）1900×900×430mm 售卖规格：1台",
                 "国保保密柜Z168-V4分体双节四门指纹密码文件保险柜办公大型 A级国密级 高1900*宽900*深430）",
                 "国保（Guub）指纹锁保密柜Z168-V4",
                 "国保保密柜Z168-V4分体双节涉密文件保险柜办公大型 A级国密级 Z168-V4（1900*900*430）",
                 "国保（Guub）保密柜指纹密码锁Z168-V4",
                 "Guub国保 Z168-V4 保密柜 办公指纹电子密码锁柜资料柜 1900*900*430",
                 "Guub保密柜Z168-V4办公指纹电子密码锁文件柜保密认证分体档案柜V4",
                 "国保（Guub）指纹密码保密柜Z168-V4 分体单节",
                 "Guub保密柜Z168-V4办公指纹电子密码锁文件柜保密认证分体档案柜V4",
                 "康宁空气炸锅 WK-AF/KZ 抽出安全断电功能",
                 "科密/COMET 碎纸机，G-3250 金盾 450×390×790mm 售卖规格：1台",
                 "科密碎纸机大型 G-3250商务办公文件粉碎机（单次25张 持续50分钟 40L 可碎纸、信用卡、光盘）",
                 "科密（Comet） G-3250 25张 3*25mm 40L 20分钟 纸 针 卡 光盘 碎纸机 (计价单位：台) 银色",
                 "科密（Comet） G-3250大型碎纸机 长时间工业级碎纸机 40L大容量办公商用文件颗粒粉碎机 企业业务",
                 "科密（comet） G-3250大型商用办公碎纸机专业销毁（单次25张纸 持续50分钟 碎卡、光盘）",
                 "科密comet 金盾G-3250颗粒粉碎机碎纸机40L大容量办公商用文件大型长时间工业级碎纸机",
                 "科密(comet)金盾G-3250碎纸机 大型长时间工业级碎纸机 40L大容量办公商用文件颗粒粉碎机",
                 "科密 G-3250 大型商用办公碎纸机 专业销毁文件粉碎机（单次25张 持续50分钟 45L 可碎卡、光盘）",
                 ]
for sku_name in sku_name_list:
    start = time.time()
    result = ner(sku_name)
    end = time.time()
    print(round(end - start, 2), sku_name, "==>>", result.strip())

使用2080ti 22g显卡推理，推理过程显存占用约为12G，结果输出：

1.68 齐心(Comix) 10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 60mm 可装600张 老款 ==>> {"品牌": "齐心", "型号": "AG-60", "主商品": "文件资料盒"}
1.13 10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 60mm 可装600张 ==>> {"品牌": "AG", "型号": "AG-60", "主商品": "文件资料盒"}
1.17 齐心10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 30mm 可装300张 ==>> {"品牌": "齐心", "型号": "AG-60", "主商品": "文件资料盒"}
1.56 SPC 韩国SPC移动平台打孔机，FP（filepecker）-I（X） 售卖规格：1台 ==>> {"品牌": "SPC", "型号": "FP（filepecker）-I（X）", "主商品": "移动平台打孔机"}
1.36 国保 分体指纹保密柜，Z168-V4 符合国家保密标准(BMB54-2020）1900×900×430mm 售卖规格：1台 ==>> {"品牌": "国保", "型号": "Z168-V4", "主商品": "分体指纹保密柜"}
1.43 国保保密柜Z168-V4分体双节四门指纹密码文件保险柜办公大型 A级国密级 高1900*宽900*深430） ==>> {"品牌": "国保", "型号": "Z168-V4", "主商品": "四门指纹密码文件保险柜"}
1.28 国保（Guub）指纹锁保密柜Z168-V4 ==>> {"品牌": "国保", "型号": "Z168-V4", "主商品": "指纹锁保密柜"}
1.22 国保保密柜Z168-V4分体双节涉密文件保险柜办公大型 A级国密级 Z168-V4（1900*900*430） ==>> {"品牌": "国保", "型号": "Z168-V4", "主商品": "保密柜"}
1.33 国保（Guub）保密柜指纹密码锁Z168-V4 ==>> {"品牌": "国保", "型号": "Z168-V4", "主商品": "保密柜指纹密码锁"}
1.3 Guub国保 Z168-V4 保密柜 办公指纹电子密码锁柜资料柜 1900*900*430 ==>> {"品牌": "Guub国保", "型号": "Z168-V4", "主商品": "保密柜"}
1.2 Guub保密柜Z168-V4办公指纹电子密码锁文件柜保密认证分体档案柜V4 ==>> {"品牌": "Guub", "型号": "Z168-V4", "主商品": "保密柜"}
1.28 国保（Guub）指纹密码保密柜Z168-V4 分体单节 ==>> {"品牌": "国保", "型号": "Z168-V4", "主商品": "指纹密码保密柜"}
1.21 Guub保密柜Z168-V4办公指纹电子密码锁文件柜保密认证分体档案柜V4 ==>> {"品牌": "Guub", "型号": "Z168-V4", "主商品": "保密柜"}
1.2 康宁空气炸锅 WK-AF/KZ 抽出安全断电功能 ==>> {"品牌": "康宁", "型号": "WK-AF/KZ", "主商品": "空气炸锅"}
1.25 科密/COMET 碎纸机，G-3250 金盾 450×390×790mm 售卖规格：1台 ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "碎纸机"}
1.35 科密碎纸机大型 G-3250商务办公文件粉碎机（单次25张 持续50分钟 40L 可碎纸、信用卡、光盘） ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "商务办公文件粉碎机"}
1.18 科密（Comet） G-3250 25张 3*25mm 40L 20分钟 纸 针 卡 光盘 碎纸机 (计价单位：台) 银色 ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "针"}
1.31 科密（Comet） G-3250大型碎纸机 长时间工业级碎纸机 40L大容量办公商用文件颗粒粉碎机 企业业务 ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "大型碎纸机"}
1.39 科密（comet） G-3250大型商用办公碎纸机专业销毁（单次25张纸 持续50分钟 碎卡、光盘） ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "大型商用办公碎纸机"}
1.21 科密comet 金盾G-3250颗粒粉碎机碎纸机40L大容量办公商用文件大型长时间工业级碎纸机 ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "粉碎机"}
1.25 科密(comet)金盾G-3250碎纸机 大型长时间工业级碎纸机 40L大容量办公商用文件颗粒粉碎机 ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "碎纸机"}
1.39 科密 G-3250 大型商用办公碎纸机 专业销毁文件粉碎机（单次25张 持续50分钟 45L 可碎卡、光盘） ==>> {"品牌": "科密", "型号": "G-3250", "主商品": "大型商用办公碎纸机"}