Supra-Router-51M-zh-cn · Multi-Task Infrastructure Routing Model

你好！请问大熊猫主要生活在中国的哪个省份？

Domain: English Language | Complexity: 2 | Math: False | Code: False | Route: small model | Justification: Automated override: Simple text or factual task with low complexity (2), perfectly optimized for an edge SLM.

请用 Python 写一段快速排序算法的代码，并简单解释它的时间复杂度。

Domain: Programming | Complexity: 2 | Math: False | Code: True | Route: big model | Justification: Automated override: Task complexity is high (2) or involves technical logic (Math: False, Code: True), requiring frontier capabilities.

如果一个等差数列的首项是3，公差是5，请问第12项是多少？让我们逐步思考。

Domain: Mathematics | Complexity: 2 | Math: True | Code: False | Route: big model | Justification: Automated override: Task complexity is high (2) or involves technical logic (Math: True, Code: False), requiring frontier capabilities.

使用

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "aifeifei798/Supra-Router-51M-zh-cn"

print("[*] Initializing local infrastructure router...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, device_map="auto"
)
model.eval()

# Example prompt showcasing keyword-trap evasion
user_prompt = "如果一个等差数列的首项是3，公差是5，请问第12项是多少？让我们逐步思考。"
print(user_prompt)
# Format to match internal SFT attention alignment
formatted_input = f"Task: {user_prompt}\nAnalysis: "
inputs = tokenizer(formatted_input, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

generated_ids = outputs[0][inputs["input_ids"].shape[1] :]
print(tokenizer.decode(generated_ids, skip_special_tokens=True).strip())

TRL 训练脚本

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig

# ================= 1. 基础配置 =================
MODEL_ID = "./Supra-Router-51M"
DATA_PATH = (
    "./Prompt-Routing-Dataset-zh-cn/router_dataset_openai_gpt-oss-120b-zhcn.jsonl"
)
OUTPUT_DIR = "./Supra-Router-51M-zh"

print("[*] 正在加载 Tokenizer 和 模型...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16
)

print(f"[*] 正在加载中文数据集: {DATA_PATH}")
dataset = load_dataset("json", data_files=DATA_PATH)["train"]


# ================= 2. 数据格式化与双列拆分 =================
def format_prompts(batch):
    prompts = []
    completions = []
    for i in range(len(batch["prompt"])):
        raw_prompt = batch["prompt"][i]

        if len(raw_prompt) > 1500:
            raw_prompt = raw_prompt[:1500] + "\n...[Text Truncated]..."

        # 1. 输入部分：去掉冒号后面的空格，防止 BPE 分词器边界混淆
        task_input = f"Task: {raw_prompt}\nAnalysis:"

        # 2. 输出部分：在开头 Domain 前面主动加一个空格，确保无缝对齐
        output_label = (
            f" Domain: {batch['primary_domain'][i]} | "
            f"Complexity: {batch['complexity_score'][i]} | "
            f"Math: {batch['math_task'][i]} | "
            f"Code: {batch['coding_task'][i]} | "
            f"Route: {batch['routing_choice'][i]} | "
            f"Justification: {batch['routing_justification'][i]}"
            f"{tokenizer.eos_token}"
        )

        prompts.append(task_input)
        completions.append(output_label)

    return {"prompt": prompts, "completion": completions}


print("[*] 正在格式化数据集并进行双列映射...")
dataset = dataset.map(format_prompts, batched=True, remove_columns=dataset.column_names)

# ================= 3. 训练参数配置 =================
training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    max_length=2048,
    completion_only_loss=True,  # 开启原生答案专注模式
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    learning_rate=3e-5,
    logging_steps=10,
    num_train_epochs=4,
    save_strategy="epoch",
    bf16=True,
    optim="adamw_torch",
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=training_args,
    processing_class=tokenizer,
)

# ================= 4. 启动微调 =================
print("[*] 万事俱备，开始以新版原生【答案专注模式】进行微调！")
trainer.train()

print(f"[*] 训练大功告成！正在保存至 {OUTPUT_DIR}")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print("[[*] 全部任务结束。")

completion_only_loss

completion_only_loss（仅对答案计算损失）是监督微调（SFT）和对齐训练（Alignment）中非常核心、甚至可以说是必用的一项技术。

为了让你彻底明白它的底层逻辑，我们从“普通训练”和“答案专注训练”的对比、PyTorch 的底层实现、以及为什么它对小模型至关重要这三个维度来剖析。

---

一、 核心对比：普通 SFT vs Completion-Only SFT

在因果语言模型（Causal LM）的日常训练中，模型的任务是预测下一个 Token。

假设一条完整的训练数据拼接后是：

Task: 1+1等于几？\nAnalysis: 等于2。

1. 普通的 SFT 训练（Full Sequence Loss）

模型需要对每一个 Token 都计算预测损失（Loss），并进行梯度回传：

• 模型预测 Task 后面的第一个字（“1”） -> 计算 Loss
• 模型预测 “1” 后面的 “+” -> 计算 Loss
• ……
• 模型预测 Analysis: 后面的 “等” -> 计算 Loss
• 模型预测 “等” 后面的 “于” -> 计算 Loss

弊端：模型在训练时，不仅要学习怎么回答“等于2”，还要花大量的脑力去学习和记忆“1+1等于几？”这段提示词。但在实际使用时，提示词是用户输入的，模型根本不需要去预测它。

2. Completion-Only SFT 训练

模型依然可以看见完整的 Task: 1+1等于几？\nAnalysis: 等于2。（用于提供完整的上下文注意力），但在计算损失时：

• 提示词（Prompt）部分：不计算 Loss，完全忽略。
• 回答（Completion）部分：正常计算 Loss。

模型唯一的优化目标，就是在看到前面的输入后，如何精准地吐出后面的答案。

---

二、 底层原理：PyTorch 的 -100 魔法

在底层代码中，这个机制是通过 PyTorch 交叉熵损失函数（CrossEntropyLoss）中的一个特殊参数 ignore_index 来实现的。

在 PyTorch 和 Hugging Face Transformers 中，ignore_index 的默认值是 -100。任何标签值被设为 -100 的 Token，在计算 Loss 时都会被完全忽略，梯度不会回传。

具体的 Token 与 Label 变化：

假设我们要训练一句话，分词后的 input_ids 如下。我们在后台为它准备了对应的 labels：

| 文本内容 | Task: | 大熊猫 | 住哪？ | \nAnalysis: | Domain: | Mathematics | | |

| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |

| input_ids (输入) | 1203 | 3042 | 9842 | 4322 | 567 | 8831 | 102 |

| labels (标签) | -100 | -100 | -100 | -100 | 567 | 8831 | 102 |

1. 在 Forward（前向传播）时，模型能看到完整的 input_ids（包含大熊猫、住哪等上下文）。
2. 在计算 Loss 时，PyTorch 看到前四个 Label 是 -100，就会直接跳过，只对后三个真正的标签值计算 CrossEntropy。
3. 这样，模型的所有参数权重，都只会为了“在看到前四个 Token 时，如何更准地预测出后三个 Token”而发生改变。

---

三、 为什么新版 TRL 的“双列自动检测”如此方便？

在过去（比如旧版的 trl 库中），我们需要使用 DataCollatorForCompletionOnlyLM。 它的工作原理是：

1. 把整段文本拼成一个大 text。
2. 在分词后的 Token 序列里，拿着一根“放大镜”，去人肉搜索 \nAnalysis: 这段字符串对应的 Token ID 边界。
3. 把边界之前的所有 Label 强行改成 -100。

缺点：非常容易因为分词器（Tokenizer）在特殊字符或空格处的切词变化（比如我们遇到的 Mismatch 问题），导致“放大镜”找不到边界，进而导致整条数据都无法计算 Loss，或者报错。

新版 TRL 的“双列原生检测”：

当你在 dataset 里保留 prompt 和 completion 两个独立的列，并在 SFTConfig 中开启 completion_only_loss=True 时，TRL 采取了更聪明的做法：

1. 它先单独对 prompt 列进行分词，得到它的 Token 长度（比如是 L）。
2. 它再对 prompt + completion 整体进行分词，得到一整个 Token 序列。
3. 它直接将这个完整序列中，前 L 个位置的 labels 暴力抹成 -100。

这是一种纯粹数学/数组切片上的操作，不再依赖于字符串的模式匹配。不仅速度极快，而且只要做好了空格对齐，就绝对不会出现对齐失败的警告。

---

四、 总结：它对 51M 小模型带来的质变

这就是为什么你在微调 51M 模型时，开启这个功能后，指标发生了翻天覆地的变化：

1. 参数解放：51M 模型太小了，如果让它去背诵你的 992 条中文 Prompt（其中有些可能长达几千字），它会因为记忆载荷过大而发生“模式坍塌”。开启后，它只需要背诵极其简短、高度重复的英文路由标签。
2. 收敛极快：因为预测目标变成了极度简单的分类选择（Domain, Code, Math 等），这相当于把一个复杂的文本生成任务（Generation）降维成了简单的选择分类任务（Classification）。这也是为什么你的 Loss 在第二轮就降到了 0.1 附近，准确率直接高达 97%+。
3. 泛化保留：因为模型不需要在训练中改变它的“中文语言表征”（因为它不用预测中文），它在微调后，依然完美保留了预训练时学到的常识和逻辑能力，从而在推理时能听懂“大熊猫”和“等差数列”的区别。

Supra-Router-51M · Multi-Task Infrastructure Routing Model

About the Model

GGUF model here

Supra-Router-51M is an ultra-lightweight, high-speed infrastructure traffic controller optimized for localized edge orchestration. With only 51.7 million parameters, this micro-LLM acts as a defensive gateway for multi-model ecosystems, accurately determining when user requests can be processed locally by an Edge SLM or when they must be triaged to a cloud-hosted frontier intelligence layer.

The model was built by fine-tuning a pre-trained 51M base on the SupraLabs/Prompt-Routing-Dataset (992 rows). Rather than acting as a naive binary classifier, the model uses Multi-Task Sequence Generation to map out the underlying properties of a prompt before predicting the final routing token, anchoring its attention heads to robust language and structural logic features.

---

Multi-Task Decision Sequence

To run inference, wrap your user query inside the structural framing tokens used during training (Task: [Prompt]\nAnalysis: ). The model will output a deterministic, pipe-separated string containing the full telemetry of the prompt's cognitive requirements:

Expected Output Target Schema:

Domain: [Semantic Field] | Complexity: [1-5] | Math: [True/False] | Code: [True/False] | Route: [small model/big model] | Justification: [Rule-driven infrastructure reasoning]

Why this works:

By forcing a sub-100M parameter model to calculate the semantic domain, structural complexity, and technical flags before it emits the final Route token, the network effectively runs an internal feature-activation map. This multi-task sequence prevents localized weight collapse and guarantees stable routing boundaries.

Training Telemetry & Optimization

• Dataset Source: SupraLabs/Prompt-Routing-Dataset (992 samples)
• Training Duration: 5 Epochs
• Checkpoint Selection: Peak generalization was reached during Epoch 3 (eval_loss: 0.1342). To eliminate late-stage micro-model memorization and validation drift, the training state was automatically rewound and saved at this numerical peak.
• Precision: bfloat16
• Hardware Footprint: Optimized sequence processing length of 3840 tokens, ensuring rapid inference execution with negligible CPU/GPU overhead (sub-millisecond generation speeds).

Inference & Gateway Implementation

Use this direct script to test or wrap the model inside a live production orchestrator or FastAPI gateway. It enforces greedy decoding (do_sample=False) for maximum decision stability.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "aifeifei798/Supra-Router-51M-zh-cn"

print("[*] Initializing local infrastructure router...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    device_map="auto"
)
model.eval()

# Example prompt showcasing keyword-trap evasion
user_prompt = "你好！请问大熊猫主要生活在中国的哪个省份？"

# Format to match internal SFT attention alignment
formatted_input = f"Task: {user_prompt}\nAnalysis: "
inputs = tokenizer(formatted_input, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False, 
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
print(tokenizer.decode(generated_ids, skip_special_tokens=True).strip())

Proven Benchmarks & Defensive Boundaries

During edge validation testing, Supra-Router-51M demonstrated robust resilience against adversarial prompt strings:

• Keyword Trap Evasion: Successfully identifies semantic context rather than matching tokens. Prompts containing words like "script" or "calculus" are correctly parsed as creative writing (not programming/math code) and routed locally to the small model when complexity is low.
• Complexity-Driven Safety Net: In instances where programming syntax or technical boundaries are ambiguous (e.g., complex regex or architectural database frames), the model naturally scales its evaluation metrics to Complexity: 3, automatically triggering a big model route override.
• Deterministic Offloading: Safely captures multi-step logic paths, calculus concepts, and code generation scripts, instantly assigning them to cloud-scale frontier endpoints.