laya-ad / code /scripts /clef_teacher.py
laya-ad v1.1：六精度权重 + 模型卡 + 项目代码快照
71dabc7 verified 9 days ago
26.4 kB
#!/usr/bin/env python3
"""Clef-Flash teacher 接入（RTX 5060 Ti 16GB · 4-bit nf4）。
模型入口是仓库自带的 ``models/clef-flash/joint_schema_model.py``（Qwen3.5-9B 视觉骨干 +
joint schema head）。脚本提供两条路径：
1. 用 bnb nf4 4-bit 在**单卡**上加载 teacher，并把 joint head 留在 bf16；
2. ``--smoke``：一张 224x224 随机图 + 一句中文 state，问一道 choice 题，打印概率后 exit 0；
3. ``--label-all``：流式读 ``data/processed/channel/messages.jsonl``，把每条文本（及其已落盘
   图片）交给 teacher，问一道 ``{"ad","normal"}`` 的 choice 题，写
   ``data/processed/teacher/channel_labels.jsonl``；每 500 条 flush 一次，重跑按 message_id 续跑。
用法::
    .venv/Scripts/python.exe scripts/clef_teacher.py --smoke
    .venv/Scripts/python.exe scripts/clef_teacher.py --label-all data/processed/channel/messages.jsonl
为什么 lm_head 必须跳过量化
---------------------------
``JointSchemaHead.forward`` 用 backbone 的输出词嵌入做 lexical 特征：
``output_embedding_weight[token_ids].mean(dim=0)``（joint_schema_model.py:382）。
bnb 4-bit 的 ``Params4bit`` 不支持按 token id 索引，量化后这条路径会直接报错/取到错值，
所以 ``llm_int8_skip_modules=["lm_head"]`` 让 lm_head 保持 bf16。
"""

from __future__ import annotations

import os

# 必须在第一次 CUDA 分配之前设置：expandable_segments 降低长序列激活的碎片化，
# 16GB 卡上 4-bit 权重（~5GB）之后剩余空间由激活独占。Windows 分配器不支持该选项。
if os.name != "nt":
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import argparse
import contextlib
import json
import sys
import time
from pathlib import Path
from typing import Any

import torch

REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MODEL_DIR = REPO_ROOT / "models" / "clef-flash"
DEFAULT_MESSAGES = REPO_ROOT / "data" / "processed" / "channel" / "messages.jsonl"
DEFAULT_LABELS = REPO_ROOT / "data" / "processed" / "teacher" / "channel_labels.jsonl"

# 602 条消息文本为空（其中 546 条连图片都没有）。空 state 会让 encode_record 产不出
# state token，这里给一个显式占位串，保持"无文本"这件事本身可见。
EMPTY_TEXT_STATE = "（无文本）"
LABEL_QUESTION_ID = "is_ad"
LABEL_INSTRUCTIONS = "这条频道帖子的内容是在做广告或引流推广吗？"
LABEL_CRITERIA = {
    "ad": "帖子在推广、引流、售卖商品，或诱导关注其他账号/平台",
    "normal": "帖子是正常的内容分享，没有广告或引流意图",
}

# 加载完成后至少要留出的空闲显存（GB）。低于该值说明激活空间不够，smoke 直接拒绝跑。
MIN_FREE_VRAM_GB = 2.0
# 加载前的最低空闲显存门槛（GB）：4-bit 权重 + 视觉塔 + head + 激活的最保守估计。
MIN_START_VRAM_GB = 7.0


def _gib(value: int | float) -> float:
    return float(value) / 1024**3


def _vram_line(tag: str, device: torch.device) -> float:
    free, total = torch.cuda.mem_get_info(device)
    allocated = torch.cuda.memory_allocated(device)
    reserved = torch.cuda.memory_reserved(device)
    print(
        f"[vram:{tag}] total={_gib(total):.2f}GB allocated={_gib(allocated):.2f}GB "
        f"reserved={_gib(reserved):.2f}GB free={_gib(free):.2f}GB"
    )
    return _gib(free)


def _resolve_device(device: str) -> torch.device:
    parsed = torch.device(device)
    if parsed.type != "cuda":
        raise SystemExit(f"--device 必须是单张 CUDA 卡（收到 {device!r}）；本脚本不做 CPU/多卡路径")
    if not torch.cuda.is_available():
        raise SystemExit("torch.cuda.is_available() == False：当前解释器没有可用的 CUDA 设备")
    index = 0 if parsed.index is None else parsed.index
    if index >= torch.cuda.device_count():
        raise SystemExit(f"--device={device} 超出可见卡数 {torch.cuda.device_count()}")
    return torch.device(f"cuda:{index}")


def _import_release_module(model_dir: Path):
    """按仓库约定把模型目录插进 sys.path 再 import 发布包里的入口函数。"""
    if not (model_dir / "joint_schema_model.py").is_file():
        raise SystemExit(f"找不到模型入口: {model_dir / 'joint_schema_model.py'}")
    entry = str(model_dir)
    if entry not in sys.path:
        sys.path.insert(0, entry)
    import joint_schema_model as release

    for name in ("load_release_model", "encode_record", "collate_records"):
        if not hasattr(release, name):
            raise SystemExit(f"{model_dir / 'joint_schema_model.py'} 缺少 {name}()")
    return release


def _build_image_only_processor(model_dir: Path):
    """构造只含 image processor + tokenizer 的 Qwen3VLProcessor。
    transformers 5.20 的 ``AutoProcessor.from_pretrained`` 会无条件实例化
    ``Qwen3VLVideoProcessor``，而该类模块顶部 ``import torchvision``；本机没装
    torchvision 时整条加载链会 ImportError。smoke 只走图像路径，因此这里显式组一个
    video_processor 缺省实现：真正传入 videos 时立即报错，绝不静默降级。
    """
    from transformers import AutoImageProcessor, AutoTokenizer
    from transformers.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor
    from transformers.video_processing_utils import VideoProcessorMixin

    # transformers 5.20.2 的 ProcessorMixin.check_argument_for_proper_class 用
    # get_possibly_dynamic_module("VideoProcessorMixin") 解析基类，但该名字没有导出到
    # transformers 顶层，回退路径又会踩 `_extra_content` 缺失的 AttributeError。
    # 这里把类补到顶层，让 isinstance 校验能正常通过（不改库文件）。
    import transformers

    if not hasattr(transformers, "VideoProcessorMixin"):
        transformers.VideoProcessorMixin = VideoProcessorMixin

    class _UnavailableVideoProcessor(VideoProcessorMixin):
        model_input_names = ["pixel_values_videos"]

        def __call__(self, *_args, **_kwargs):
            raise RuntimeError(
                "torchvision 未安装：本脚本只接入图像输入，视频路径不可用"
                "（安装 torchvision 后重跑即可使用官方 processor）"
            )

    image_processor = AutoImageProcessor.from_pretrained(str(model_dir))
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
    return Qwen3VLProcessor(
        image_processor=image_processor,
        tokenizer=tokenizer,
        video_processor=_UnavailableVideoProcessor(),
    )


@contextlib.contextmanager
def _patched_auto_processor(processor):
    """让 ``load_release_model`` 内部的 AutoProcessor.from_pretrained 返回现成 processor。"""
    from transformers import AutoProcessor

    saved = AutoProcessor.__dict__.get("from_pretrained")
    AutoProcessor.from_pretrained = classmethod(lambda cls, *args, **kwargs: processor)
    try:
        yield
    finally:
        if saved is None:
            del AutoProcessor.from_pretrained
        else:
            setattr(AutoProcessor, "from_pretrained", saved)


def _load_processor(model_dir: Path) -> tuple[object, bool, str]:
    """优先用官方 AutoProcessor；缺 torchvision 时退化为 image-only 构造。"""
    from transformers import AutoProcessor

    try:
        processor = AutoProcessor.from_pretrained(str(model_dir))
        return processor, False, "official"
    except ImportError as error:
        if "torchvision" not in str(error).lower():
            raise
        return _build_image_only_processor(model_dir), True, "image_only(torchvision missing)"


def _build_quant_config(dtype: torch.dtype):
    from transformers import BitsAndBytesConfig

    return BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=dtype,
        bnb_4bit_use_double_quant=True,
        # 见模块 docstring：joint head 拿 lm_head 权重按 token id 查表，必须保持浮点。
        llm_int8_skip_modules=["lm_head"],
    )


def verify_quantization(model, dtype: torch.dtype = torch.bfloat16) -> int:
    """确认真的走了 4-bit，且 lm_head/head 没被误量化。返回 Linear4bit 层数。"""
    from bitsandbytes.nn import Linear4bit, Params4bit

    backbone = model.language_model
    quantized_layers = sum(1 for module in backbone.modules() if isinstance(module, Linear4bit))
    if quantized_layers == 0:
        raise SystemExit("没有任何 Linear 被量化成 4-bit：检查 bitsandbytes 与 transformers 版本")

    lm_head = backbone.get_output_embeddings()
    weight = lm_head.weight
    if isinstance(weight, Params4bit):
        raise SystemExit(
            "lm_head 被 4-bit 量化了：joint head 的 lexical 路径会失败（joint_schema_model.py:382）"
        )

    head_dtype = next(model.head.parameters()).dtype
    sample_layer = next(m for m in backbone.modules() if isinstance(m, Linear4bit))
    quant_state = getattr(sample_layer.weight, "quant_state", None)
    quant_type = getattr(quant_state, "quant_type", None)
    print(
        f"[load] Linear4bit={quantized_layers} quant_type={quant_type} "
        f"lm_head_dtype={weight.dtype} head_dtype={head_dtype} "
        f"head_device={next(model.head.parameters()).device}"
    )
    if str(quant_type).lower() != "nf4":
        raise SystemExit(f"4-bit 量化类型不是 nf4（实际 {quant_type!r}），检查 BitsAndBytesConfig")
    if head_dtype != dtype:
        raise SystemExit(f"joint head 期望 {dtype}，实际 {head_dtype}")
    return quantized_layers


def load_teacher(
    model_dir: Path,
    device: str = "cuda:0",
    dtype: torch.dtype = torch.bfloat16,
) -> tuple[object, object, object]:
    """4-bit nf4 加载 teacher，返回 (release 模块, model, processor)。"""
    target = _resolve_device(device)
    torch.cuda.set_device(target)

    free_before = _gib(torch.cuda.mem_get_info(target)[0])
    print(f"[load] device={target} free_before={free_before:.2f}GB model_dir={model_dir}")
    if free_before < MIN_START_VRAM_GB:
        raise SystemExit(
            f"加载前空闲显存只有 {free_before:.2f}GB，低于 {MIN_START_VRAM_GB}GB 门槛；"
            "先释放占用显存的进程再试"
        )

    release = _import_release_module(model_dir)
    processor, use_fallback, note = _load_processor(model_dir)
    print(f"[load] processor={type(processor).__name__} source={note}")
    if use_fallback:
        with _patched_auto_processor(processor):
            model, loaded_processor = release.load_release_model(
                model_dir,
                device=str(target),
                dtype=dtype,
                quantization_config=_build_quant_config(dtype),
            )
        if loaded_processor is not processor:
            raise SystemExit(
                f"processor 注入失败：load_release_model 返回了 {type(loaded_processor).__name__}"
            )
    else:
        model, processor = release.load_release_model(
            model_dir,
            device=str(target),
            dtype=dtype,
            quantization_config=_build_quant_config(dtype),
        )
    model.eval()
    verify_quantization(model, dtype)

    free_after = _vram_line("after_load", target)
    if free_after < MIN_FREE_VRAM_GB:
        raise SystemExit(
            f"加载后只剩 {free_after:.2f}GB 空闲显存，低于激活所需的 {MIN_FREE_VRAM_GB}GB 余量"
        )
    return release, model, processor


def build_smoke_record(seed: int) -> tuple[dict, object]:
    """一张 224x224 随机图 + 一句中文 state + 一道三选一 choice 题。"""
    import numpy as np
    from PIL import Image

    rng = np.random.default_rng(seed)
    pixels = rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8)
    image = Image.fromarray(pixels, mode="RGB")

    record = {
        "id": "smoke",
        "state": "客服工单：用户反馈订单 A1024 在支付页面连续三次失败，要求今天内给出处理方案。",
        "images": [image],
        "questions": {
            "team": {
                "type": "choice",
                "instructions": "根据状态和图片，这张工单应该交给哪个团队处理？",
                "criteria": {
                    "billing": "支付、退款、账单相关的问题",
                    "technical": "程序缺陷、服务中断或报错",
                    "general": "一般咨询、功能介绍或账户操作",
                },
            }
        },
    }
    return record, image


def run_smoke(
    release,
    model,
    processor,
    max_length: int = 4096,
    seed: int = 0,
) -> int:
    device = next(model.parameters()).device
    record, image = build_smoke_record(seed)
    print(f"[smoke] image={image.size[0]}x{image.size[1]} mode={image.mode} seed={seed}")

    encoded = release.encode_record(
        processor.tokenizer,
        record,
        max_length=max_length,
        processor=processor,
    )
    print(
        f"[smoke] input_tokens={len(encoded.input_ids)} questions={len(encoded.questions)} "
        f"media={sorted((encoded.media or {}).keys())}"
    )

    batch = release.collate_records([encoded], processor.tokenizer.pad_token_id, device)
    with torch.inference_mode():
        logits = model(batch)[0]

    exit_code = 0
    for question, question_logits in zip(encoded.questions, logits):
        probabilities = question_logits.float().softmax(-1)
        if not torch.isfinite(probabilities).all():
            print(f"[smoke] FAIL: {question.question_id} 概率含 NaN/Inf", file=sys.stderr)
            exit_code = 1
            continue

        pairs = sorted(zip(question.option_ids, probabilities.tolist()), key=lambda kv: -kv[1])
        print(f"[smoke] question={question.question_id} options={len(pairs)}")
        for option_id, probability in pairs:
            print(f"  {option_id:>10s}  p={probability:.6f}")
        best_option, best_probability = pairs[0]
        total = sum(probability for _, probability in pairs)
        print(
            f"[smoke] argmax={best_option} p={best_probability:.6f} sum={total:.6f} "
            f"logits={[round(x, 4) for x in question_logits.float().tolist()]}"
        )
        if abs(total - 1.0) > 1e-4:
            print(f"[smoke] FAIL: 概率和 {total:.6f} 偏离 1", file=sys.stderr)
            exit_code = 1

    _vram_line("after_smoke", device)
    print(f"[smoke] peak_allocated={_gib(torch.cuda.max_memory_allocated(device)):.2f}GB")
    return exit_code


def _count_lines(path: Path) -> int:
    with path.open("rb") as handle:
        return sum(1 for _ in handle)


def _load_done_ids(path: Path) -> set[str]:
    """读已有输出里的 message_id；顺手截掉崩溃留下的半行。"""
    if not path.exists():
        return set()
    raw = path.read_text(encoding="utf-8")
    if raw and not raw.endswith("\n"):
        cut = raw.rfind("\n") + 1
        path.write_text(raw[:cut], encoding="utf-8")
        raw = raw[:cut]
    done: set[str] = set()
    for line in raw.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            done.add(str(json.loads(line)["message_id"]))
        except (ValueError, KeyError):
            continue
    return done


def _iter_messages(path: Path):
    """逐行流式读取，避免把 10k 条消息一次性放进内存。"""
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                yield json.loads(line)


def _build_label_record(row: dict[str, Any]) -> tuple[dict[str, Any], int, bool]:
    """消息 → Clef record。返回 (record, 成功读入的图片数, 是否声明了图片)。"""
    from PIL import Image

    text = (row.get("text") or "").strip()
    paths = [str(path) for path in (row.get("images") or [])]
    images = []
    for path in paths:
        with Image.open(path) as image:  # copy 出来，避免文件句柄跟着 PIL 对象悬空
            images.append(image.convert("RGB").copy())
    record = {
        "id": str(row.get("message_id")),
        "state": text or EMPTY_TEXT_STATE,
        "images": images,
        "questions": {
            LABEL_QUESTION_ID: {
                "type": "choice",
                "instructions": LABEL_INSTRUCTIONS,
                "criteria": dict(LABEL_CRITERIA),
            }
        },
    }
    return record, len(images), bool(paths)


def _error_row(row: dict[str, Any], has_image: bool, error: str) -> dict[str, Any]:
    return {
        "message_id": row.get("message_id"),
        "has_image": has_image,
        "error": error,
        "date": row.get("date"),
    }


def _forward_encoded(release, model, processor, encoded, device) -> list[dict[str, float]]:
    batch = release.collate_records(encoded, processor.tokenizer.pad_token_id, device)
    with torch.inference_mode():
        logits = model(batch)
    probabilities: list[dict[str, float]] = []
    for encoding, record_logits in zip(encoded, logits):
        per_question: dict[str, float] = {}
        for question, question_logits in zip(encoding.questions, record_logits):
            values = question_logits.float().softmax(-1).tolist()
            per_question.update(zip(question.option_ids, values))
        probabilities.append(per_question)
    return probabilities


def _forward_encoded_safe(release, model, processor, encoded, device) -> list[tuple[dict | None, str | None]]:
    """单批前向；OOM / 批级异常时二分重试，把坏样本隔离到单条。"""
    try:
        return [(probabilities, None) for probabilities in _forward_encoded(release, model, processor, encoded, device)]
    except torch.OutOfMemoryError as error:
        torch.cuda.empty_cache()
        if len(encoded) == 1:
            return [(None, f"cuda_oom: {error}")]
    except Exception as error:  # noqa: BLE001 - 批级失败先二分，单条再抛出
        if len(encoded) == 1:
            return [(None, f"{type(error).__name__}: {error}")]
    midpoint = max(1, len(encoded) // 2)
    return _forward_encoded_safe(
        release, model, processor, encoded[:midpoint], device
    ) + _forward_encoded_safe(release, model, processor, encoded[midpoint:], device)


def run_label_all(release, model, processor, args) -> int:
    """批量打标：messages.jsonl → channel_labels.jsonl（流式 + 每 500 条 flush + 断点续跑）。"""
    input_path = Path(args.label_all).expanduser().resolve()
    output_path = Path(args.output).expanduser().resolve()
    errors_path = output_path.with_suffix(".errors.jsonl")
    if not input_path.is_file():
        raise SystemExit(f"--label-all 输入不存在: {input_path}")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    device = next(model.parameters()).device
    total = _count_lines(input_path)
    done = set() if args.no_resume else _load_done_ids(output_path)
    target = max(0, total - len(done))
    if args.limit is not None:
        target = min(target, args.limit)
    print(
        f"[label-all] input={input_path} lines={total} skip_done={len(done)} target={target} "
        f"output={output_path} batch={args.batch_size} flush_every={args.flush_every}"
    )

    rows: list[dict] = []
    error_rows: list[dict] = []
    counters = {"written": 0, "processed": 0, "errors": 0, "skipped": 0}
    started = time.monotonic()

    def flush():
        if rows:
            with output_path.open("a", encoding="utf-8") as handle:
                handle.write("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows))
                handle.flush()
                os.fsync(handle.fileno())
            counters["written"] += len(rows)
            rows.clear()
        if error_rows:
            with errors_path.open("a", encoding="utf-8") as handle:
                handle.write("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in error_rows))
                handle.flush()
                os.fsync(handle.fileno())
            error_rows.clear()
        elapsed = time.monotonic() - started
        rate = counters["processed"] / elapsed if elapsed > 0 else 0.0
        remaining = max(0, target - counters["processed"]) / rate / 60 if rate > 0 else float("nan")
        print(
            f"[label-all] {counters['processed']}/{target} written={counters['written']} "
            f"errors={counters['errors']} {rate:.2f} rows/s elapsed={elapsed / 60:.1f}m eta={remaining:.1f}m"
        )

    def process_batch(items: list[tuple]) -> None:
        encoded, kept = [], []
        for row, record, has_image, images_used in items:
            try:
                encoding = release.encode_record(
                    processor.tokenizer,
                    record,
                    max_length=args.max_length,
                    processor=processor,
                )
            except Exception as error:  # noqa: BLE001 - 坏样本进 errors，不中断整批
                counters["errors"] += 1
                error_rows.append(_error_row(row, has_image, f"encode {type(error).__name__}: {error}"))
                continue
            encoded.append(encoding)
            kept.append((row, has_image, images_used))

        if encoded:
            for (row, has_image, images_used), (option_probs, error) in zip(
                kept, _forward_encoded_safe(release, model, processor, encoded, device)
            ):
                if option_probs is None:
                    counters["errors"] += 1
                    error_rows.append(_error_row(row, has_image, error or "forward failed"))
                    continue
                p_ad = float(option_probs["ad"])
                p_normal = float(option_probs["normal"])
                rows.append(
                    {
                        "message_id": row.get("message_id"),
                        "p_ad": round(p_ad, 6),
                        "label": "ad" if p_ad >= p_normal else "normal",
                        "has_image": has_image,
                        "p_normal": round(p_normal, 6),
                        "images_used": images_used,
                        "date": row.get("date"),
                        "iso_time": row.get("iso_time"),
                    }
                )

    batch_items: list[tuple] = []
    for row in _iter_messages(input_path):
        message_id = str(row.get("message_id"))
        if message_id in done:
            counters["skipped"] += 1
            continue
        if args.limit is not None and counters["processed"] + len(batch_items) >= args.limit:
            break
        try:
            record, images_used, has_image = _build_label_record(row)
        except Exception as error:  # noqa: BLE001 - 单条图片/构造失败进 errors
            counters["errors"] += 1
            counters["processed"] += 1
            error_rows.append(_error_row(row, bool(row.get("images")), f"build {type(error).__name__}: {error}"))
            continue

        batch_items.append((row, record, has_image, images_used))
        if len(batch_items) >= args.batch_size:
            process_batch(batch_items)
            counters["processed"] += len(batch_items)
            batch_items.clear()
            if len(rows) + len(error_rows) >= args.flush_every:
                flush()

    if batch_items:
        process_batch(batch_items)
        counters["processed"] += len(batch_items)
    flush()

    elapsed = time.monotonic() - started
    output_lines = _count_lines(output_path) if output_path.exists() else 0
    error_lines = _count_lines(errors_path) if errors_path.exists() else 0
    print(
        f"[label-all] DONE processed={counters['processed']} written_this_run={counters['written']} "
        f"skip_done={counters['skipped']} errors={counters['errors']} elapsed={elapsed / 60:.1f}m"
    )
    print(f"[label-all] output={output_path} lines={output_lines} error_file={errors_path} error_lines={error_lines}")
    return 0 if counters["errors"] == 0 else 1



def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Clef-Flash teacher：4-bit nf4 单卡加载 + smoke 自检",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=DEFAULT_MODEL_DIR,
        help=f"模型目录（默认 {DEFAULT_MODEL_DIR}）",
    )
    parser.add_argument("--device", default="cuda:0", help="单卡设备，默认 cuda:0")
    parser.add_argument(
        "--max-length",
        type=int,
        default=4096,
        help="encode_record 的 token 上限（默认 4096，给 16GB 卡的激活留余量）",
    )
    parser.add_argument("--seed", type=int, default=0, help="随机图种子（默认 0）")

    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--smoke", action="store_true", help="跑一张随机图 + 中文 state 的 choice 题")
    mode.add_argument(
        "--label-all",
        metavar="INPUT",
        default=None,
        help=f"批量打标：读 messages.jsonl（默认 {DEFAULT_MESSAGES}），写 channel_labels.jsonl",
    )

    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_LABELS,
        help=f"--label-all 输出路径（默认 {DEFAULT_LABELS}）",
    )
    parser.add_argument("--batch-size", type=int, default=8, help="每次前向的记录数（默认 8）")
    parser.add_argument("--flush-every", type=int, default=500, help="每多少条落盘一次（默认 500）")
    parser.add_argument("--limit", type=int, default=None, help="只处理前 N 条（调试用，默认全部）")
    parser.add_argument(
        "--no-resume",
        action="store_true",
        help="不续跑：忽略输出文件里已有的 message_id（注意输出仍是追加写，需先自行删除旧文件）",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    model_dir = args.model_dir.expanduser().resolve()

    release, model, processor = load_teacher(model_dir, device=args.device)
    if args.smoke:
        return run_smoke(release, model, processor, max_length=args.max_length, seed=args.seed)
    return run_label_all(release, model, processor, args)


if __name__ == "__main__":
    sys.exit(main())