# 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-Instruct"

# 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]

    return response


sku_name_list = ["大嘴猴（paul frank)花漾柔护洗衣液",
                 "得力（deli）珊瑚海A3打印纸 70g500张*5包一箱 双面绘图纸复印纸 手抄报画纸 整箱2500张7365【销冠系列】",
                 "优露清 多功能全能清洁剂去污除垢厨房油污不锈钢瓷砖玻璃地板清洗剂 1瓶装",
                 "滴露（Dettol）消毒液消毒水1.2L衣物除菌液家居宠物环境地板杀菌除螨 非84酒精",
                 "得力（deli）佳铂A4打印纸 70g500张 高档单包复印纸 合同标书彩打纸 打印书写3584【纸中贵族】",
                 "得力（deli）1mm加厚高档磨砂档案盒A4塑料PP文件夹文件盒 财会文件收纳资料盒资料册 加厚款35mm蓝色-岩石纹理 5605",
                 "心相印DT25150经典系列抽纸2层150抽面巾纸餐巾纸卫生纸家用厕纸擦手纸 8包",
                 "FUJIFILM富士施乐CT203435原装墨粉盒适用于施乐AP4570/5570墨盒（约印36000张）",
                 "艾贝丽家用柔音豆浆机 ABL-DJ31",
                 "得力财务装订机 凭证装订机 GB506 全自动热熔胶装 电动打孔机档案记账财务 企业业务",
                 "得力(deli)1只55mmA4文件盒批发 加厚党建档案盒 桌面考试收纳 财会用品 5603黑色",

                 ]
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())

结果输出：

text

99 齐心(Comix) 10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 60mm 可装600张 老款 ==>> {'品牌': '齐心', '型号': 'AG-60', '主商品': '牛皮纸档案盒'}
57 10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 60mm 可装600张 ==>> {'品牌': '10个装', '型号': 'AG-60', '主商品': '牛皮纸档案盒'}
49 齐心10个装 60mm高质感牛皮纸档案盒 A4文件资料盒 办公文具AG-60 30mm 可装300张 ==>> {'品牌': '齐心', '型号': 'AG-60', '主商品': '牛皮纸档案盒'}
85 SPC 韩国SPC移动平台打孔机，FP（filepecker）-I（X） 售卖规格：1台 ==>> {'品牌': 'SPC', '型号': 'FP（filepecker）-I（X）', '主商品': '移动平台打孔机'}
58 国保 分体指纹保密柜，Z168-V4 符合国家保密标准(BMB54-2020）1900×900×430mm 售卖规格：1台 ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '分体指纹保密柜'}
43 国保保密柜Z168-V4分体双节四门指纹密码文件保险柜办公大型 A级国密级 高1900*宽900*深430） ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '保密柜'}
51 国保（Guub）指纹锁保密柜Z168-V4 ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '指纹锁保密柜'}
41 国保保密柜Z168-V4分体双节涉密文件保险柜办公大型 A级国密级 Z168-V4（1900*900*430） ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '保密柜'}
39 国保（Guub）保密柜指纹密码锁Z168-V4 ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '保密柜'}
4 Guub国保 Z168-V4 保密柜 办公指纹电子密码锁柜资料柜 1900*900*430 ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '保密柜'}
35 Guub保密柜Z168-V4办公指纹电子密码锁文件柜保密认证分体档案柜V4 ==>> {'品牌': 'Guub', '型号': 'Z168-V4', '主商品': '保密柜'}
45 国保（Guub）指纹密码保密柜Z168-V4 分体单节 ==>> {'品牌': '国保', '型号': 'Z168-V4', '主商品': '指纹密码保密柜'}
39 Guub保密柜Z168-V4办公指纹电子密码锁文件柜保密认证分体档案柜V4 ==>> {'品牌': 'Guub', '型号': 'Z168-V4', '主商品': '保密柜'}
39 康宁空气炸锅 WK-AF/KZ 抽出安全断电功能 ==>> {'品牌': '康宁', '型号': 'WK-AF/KZ', '主商品': '空气炸锅'}
45 科密/COMET 碎纸机，G-3250 金盾 450×390×790mm 售卖规格：1台 ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '碎纸机'}
44 科密碎纸机大型 G-3250商务办公文件粉碎机（单次25张 持续50分钟 40L 可碎纸、信用卡、光盘） ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '碎纸机'}
45 科密（Comet） G-3250 25张 3*25mm 40L 20分钟 纸 针 卡 光盘 碎纸机 (计价单位：台) 银色 ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '碎纸机'}
44 科密（Comet） G-3250大型碎纸机 长时间工业级碎纸机 40L大容量办公商用文件颗粒粉碎机 企业业务 ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '碎纸机'}
44 科密（comet） G-3250大型商用办公碎纸机专业销毁（单次25张纸 持续50分钟 碎卡、光盘） ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '碎纸机'}
84 科密comet 金盾G-3250颗粒粉碎机碎纸机40L大容量办公商用文件大型长时间工业级碎纸机 ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '金盾G-3250颗粒粉碎机'}
44 科密(comet)金盾G-3250碎纸机 大型长时间工业级碎纸机 40L大容量办公商用文件颗粒粉碎机 ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '碎纸机'}
57 科密 G-3250 大型商用办公碎纸机 专业销毁文件粉碎机（单次25张 持续50分钟 45L 可碎卡、光盘） ==>> {'品牌': '科密', '型号': 'G-3250', '主商品': '大型商用办公碎纸机'}
3 大嘴猴（paul frank)花漾柔护洗衣液 ==>> {'品牌': '大嘴猴', '型号': None, '主商品': '花漾柔护洗衣液'}
39 得力（deli）珊瑚海A3打印纸 70g500张*5包一箱 双面绘图纸复印纸 手抄报画纸 整箱2500张7365【销冠系列】 ==>> {'品牌': '得力', '型号': '7365', '主商品': '双面绘图纸'}
22 优露清 多功能全能清洁剂去污除垢厨房油污不锈钢瓷砖玻璃地板清洗剂 1瓶装 ==>> {'品牌': '优露清', '型号': None, '主商品': '多功能全能清洁剂'}
07 滴露（Dettol）消毒液消毒水1.2L衣物除菌液家居宠物环境地板杀菌除螨 非84酒精 ==>> {'品牌': '滴露', '型号': None, '主商品': '消毒液'}
31 得力（deli）佳铂A4打印纸 70g500张 高档单包复印纸 合同标书彩打纸 打印书写3584【纸中贵族】 ==>> {'品牌': '得力', '型号': '3584', '主商品': '打印纸'}
31 得力（deli）1mm加厚高档磨砂档案盒A4塑料PP文件夹文件盒 财会文件收纳资料盒资料册 加厚款35mm蓝色-岩石纹理 5605 ==>> {'品牌': '得力', '型号': '5605', '主商品': '档案盒'}
43 心相印DT25150经典系列抽纸2层150抽面巾纸餐巾纸卫生纸家用厕纸擦手纸 8包 ==>> {'品牌': '心相印', '型号': 'DT25150', '主商品': '抽纸'}
69 FUJIFILM富士施乐CT203435原装墨粉盒适用于施乐AP4570/5570墨盒（约印36000张） ==>> {'品牌': '富士施乐', '型号': 'CT203435', '主商品': '原装墨粉盒'}
53 艾贝丽家用柔音豆浆机 ABL-DJ31 ==>> {'品牌': '艾贝丽', '型号': 'ABL-DJ31', '主商品': '家用柔音豆浆机'}
2 得力财务装订机 凭证装订机 GB506 全自动热熔胶装 电动打孔机档案记账财务 企业业务 ==>> {'品牌': '得力', '型号': None, '主商品': '财务装订机'}
29 得力(deli)1只55mmA4文件盒批发 加厚党建档案盒 桌面考试收纳 财会用品 5603黑色 ==>> {'品牌': '得力', '型号': '5603', '主商品': '文件盒'}