"""角色段掩码演示：python 10_sft_audit.py。可选训练配方见脚本末尾。"""

segments = [
    ("system", [101, 102, 103]),
    ("user", [201, 202, 203, 204]),
    ("assistant_call", [301, 302, 303]),
    ("tool", [401, 402, 403, 404]),
    ("assistant_final", [501, 502, 503]),
]

def labels(items, limit, wrong=False):
    tokens, mask = [], []
    for role, ids in items:
        train = role.startswith("assistant") or (wrong and role == "tool")
        tokens += ids
        mask += [token if train else -100 for token in ids]
    return tokens[:limit], mask[:limit]

for limit, wrong in ((30, False), (30, True), (7, False)):
    ids, y = labels(segments, limit, wrong)
    leaked = [x for x in y if 400 <= x < 500]
    supervised = sum(x != -100 for x in y)
    print(f"context={limit}, wrong_mask={wrong}, targets={supervised}, tool_observation_targets={leaked}")
assert labels(segments, 30)[1].count(-100) == 11
assert all(x == -100 for x in labels(segments, 7)[1])

# 真实 SFT 配方（本演示不执行）。先安装与 CUDA 兼容的 torch，
# 再安装 trl、transformers、datasets、peft，固定已验证版本。
# train.jsonl、dev.jsonl 每行 {"messages":[...]}，按场景家族预先切分。
# 用支持目标工具调用格式和 generation mask 的本地基座。
#
# from datasets import load_dataset
# from peft import LoraConfig
# from trl import SFTConfig, SFTTrainer
# ds = load_dataset("json", data_files={"train":"train.jsonl","validation":"dev.jsonl"})
# args = SFTConfig(output_dir="./adapter", max_length=2048,
#                  assistant_only_loss=True, packing=False,
#                  per_device_train_batch_size=1, gradient_accumulation_steps=16,
#                  learning_rate=1e-4, num_train_epochs=1, eval_strategy="steps",
#                  eval_steps=50, logging_steps=10, seed=42)
# adapter = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, task_type="CAUSAL_LM",
#                      target_modules=["q_proj","v_proj"])  # 先核对基座模块名
# trainer = SFTTrainer(model="./base_model", args=args, peft_config=adapter,
#                      train_dataset=ds["train"], eval_dataset=ds["validation"])
# 先抽样核查实际 batch 的 input_ids、labels 与聊天模板，再执行：
# trainer.train()
# trainer.save_model("./adapter/final")