Constraining Output Space for SLM Narrow Automation Optimization

pip install torch transformers accelerate

import os
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

torch.set_num_threads(os.cpu_count() or 1)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)
model.eval()

# 我们的待分类样本数据（600条记录）
LABELS = ["billing", "technical", "account"]
tickets = [
    "My card was charged twice for the same invoice.",
    "The mobile app crashes whenever I open the settings page.",
    "I need to change the email address on my profile.",
] * 200

def build_prompt(ticket):
    messages = [
        {
            "role": "system",
            "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.",
        },
        {"role": "user", "content": f"Ticket: {ticket}\nCategory:"},
    ]
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

# 此处输出未受约束，因此让模型生成简短回答并从中搜索标签（标记++）
# 每个新标记都需要单独的前向传递，每次调用处理一个工单意味着无法通过批量处理来分摊成本（时间++）
prompts = [build_prompt(t) for t in tickets]
predictions = []

# 计算推理耗时
start = time.time()

for n, prompt in enumerate(prompts, start=1):

    # 此循环在CPU上运行需要数分钟，因此需要显示进度而不是保持静默
    if n % 50 == 0:
        rate = (time.time() - start) / n
        print(f"  {n}/{len(prompts)} tickets ({rate:.2f}s each)", flush=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=8,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    # generate() 返回 prompt + 续写内容，因此在解码前需要切片去除 prompt 部分
    generated = output[0, inputs["input_ids"].shape[1] :]
    text = tokenizer.decode(generated, skip_special_tokens=True).strip().lower()

    # 与标签列表进行子字符串匹配
    predictions.append(next((label for label in LABELS if label in text), "UNPARSED"))

duration = time.time() - start

# 输出任务指标
print(f"自由形式生成耗时: {duration:.2f} 秒")
print(f"无法解析的输出: {predictions.count('UNPARSED')} / {len(predictions)}")

# 推理输出示例
for ticket, label in zip(tickets[-3:], predictions[-3:], strict=True):
    print(f"{ticket} -> {label}")

加载权重: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 263.89it/s]
  50/600 个票据 (每个 0.19 秒)
  100/600 个票据 (每个 0.19 秒)
  150/600 个票据 (每个 0.19 秒)
  200/600 个票据 (每个 0.19 秒)
  250/600 个票据 (每个 0.19 秒)
  300/600 个票据 (每个 0.19 秒)
  350/600 个票据 (每个 0.19 秒)
  400/600 个票据 (每个 0.19 秒)
  450/600 个票据 (每个 0.19 秒)
  500/600 个票据 (每个 0.19 秒)
  550/600 个票据 (每个 0.19 秒)
  600/600 个票据 (每个 0.19 秒)
自由形式生成耗时：134.01 秒
无法解析的输出：0 / 600
我的信用卡因同一发票被收取两次费用。 -> billing
每次打开设置页面时移动应用都会崩溃。 -> technical
我需要修改个人资料上的电子邮件地址。 -> technical

import os
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

torch.set_num_threads(os.cpu_count() or 1)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
model.eval()

# 我们的分类测试数据（600条记录）
LABELS = ["billing", "technical", "account"]
tickets = [
    "My card was charged twice for the same invoice.",
    "The mobile app crashes whenever I open the settings page.",
    "I need to change the email address on my profile.",
] * 200

def build_prompt(ticket):
    messages = [
        {
            "role": "system",
            "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.",
        },
        {"role": "user", "content": f"Ticket: {ticket}\nCategory:"},
    ]
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

# 提示信息以 "assistant\n" 结尾，因此模型的下一个标记将开始标签
# 只要这些第一个标记是唯一的，比较每个标签第一个标记的logits就足以确定结果
label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]
assert len(set(label_first_ids)) == len(LABELS), (
    "标签共享第一个标记；请改用完整标签序列进行评分（详见说明）。"
)
label_first_ids = torch.tensor(label_first_ids, device=model.device)

# 每个票据只需一次前向传递，无需生成循环：决策完全包含在下一个标记的logits中
prompts = [build_prompt(t) for t in tickets]
predictions = []
confidences = []

# 测量推理时间
start = time.time()
for prompt in prompts:
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        logits = model(**inputs).logits[0, -1, :]
    # 仅对标签logits进行softmax，使候选项的概率总和为1
    probs = torch.softmax(logits[label_first_ids].float(), dim=-1)
    best = int(probs.argmax())
    predictions.append(LABELS[best])
    confidences.append(float(probs[best]))
duration = time.time() - start

# 输出任务指标
print(f"受限评分耗时：{duration:.2f} 秒")
print(f"无法解析的输出：{predictions.count('UNPARSED')} / {len(predictions)}")

# 推理输出示例
for ticket, label, confidence in zip(
    tickets[-3:], predictions[-3:], confidences[-3:], strict=True
):
    print(f"{ticket} -> {label} (置信度 {confidence:.3f})")

加载权重: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 285.60it/s]
受限评分耗时：94.51 秒
无法解析的输出：0 / 600
My card was charged twice for the same invoice. -> billing (置信度 0.793)
The mobile app crashes whenever I open the settings page. -> technical (置信度 0.798)
I need to change the email address on my profile. -> technical (置信度 0.673)