SLM の出力空間制約による狭義自動化最適化手法を提案

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

# our toy data to classify (600 records)
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
    )

# nothing constrains the output here, so we let the model write a short answer and search it for a label (tokens++)
# each new token costs its own forward pass, and one ticket per call means no batching to amortize that (time++)
prompts = [build_prompt(t) for t in tickets]
predictions = []

# time inference
start = time.time()

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

    # this loop runs for minutes on CPU, so report progress rather than sitting silent
    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() returns prompt + continuation, so slice the prompt off before decoding
    generated = output[0, inputs["input_ids"].shape[1] :]
    text = tokenizer.decode(generated, skip_special_tokens=True).strip().lower()

    # substring match against the label list
    predictions.append(next((label for label in LABELS if label in text), "UNPARSED"))

duration = time.time() - start

# output task metrics
print(f"Free-form generation took: {duration:.2f} seconds")
print(f"Unparseable outputs: {predictions.count('UNPARSED')} / {len(predictions)}")

# sample of inference output
for ticket, label in zip(tickets[-3:], predictions[-3:], strict=True):
    print(f"{ticket} -> {label}")

Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 263.89it/s]
  50/600 tickets (0.19s each)
  100/600 tickets (0.19s each)
  150/600 tickets (0.19s each)
  200/600 tickets (0.19s each)
  250/600 tickets (0.19s each)
  300/600 tickets (0.19s each)
  350/600 tickets (0.19s each)
  400/600 tickets (0.19s each)
  450/600 tickets (0.19s each)
  500/600 tickets (0.19s each)
  550/600 tickets (0.19s each)
  600/600 tickets (0.19s each)
Free-form generation took: 134.01 seconds
Unparseable outputs: 0 / 600
My card was charged twice for the same invoice. -> billing
The mobile app crashes whenever I open the settings page. -> technical
I need to change the email address on my profile. -> 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()

# our toy data to classify (600 records)
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
    )

# prompt ends with "<|im_start|>assistant\n", so the model's next token starts the label
# comparing the logits of each label's FIRST token is enough to pick a winner, provided those first tokens are distinct
label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]
assert len(set(label_first_ids)) == len(LABELS), (
    "Labels share a first token; score full label sequences instead (see notes)."
)
label_first_ids = torch.tensor(label_first_ids, device=model.device)

# one forward pass per ticket, no generation loop: the decision contained entirely in the next-token logits
prompts = [build_prompt(t) for t in tickets]
predictions = []
confidences = []

# time inference
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, :]
    # softmax over just the label logits, so the probabilities sum to 1 across the candidates
    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

# output task metrics
print(f"Constrained scoring took: {duration:.2f} seconds")
print(f"Unparseable outputs: {predictions.count('UNPARSED')} / {len(predictions)}")

# sample of inference output
for ticket, label, confidence in zip(
    tickets[-3:], predictions[-3:], confidences[-3:], strict=True
):
    print(f"{ticket} -> {label} (confidence {confidence:.3f})")

Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 285.60it/s]
Constrained scoring took: 94.51 seconds
Unparseable outputs: 0 / 600
My card was charged twice for the same invoice. -> billing (confidence 0.793)
The mobile app crashes whenever I open the settings page. -> technical (confidence 0.798)
I need to change the email address on my profile. -> technical (confidence 0.673)