---
library_name: transformers
tags:
- router
- orchestrator
- slm
- edge-computing
- mixture-of-experts
- text-generation
- TRL
pipeline_tag: text-generation
model_type: llama
datasets:
- aifeifei798/Prompt-Routing-Dataset-zh-cn
language:
- en
- zh
base_model:
- SupraLabs/Supra-Router-51M
license: apache-2.0
---

<h1 align="center">Supra-Router-51M-zh-cn · Multi-Task Infrastructure Routing Model</h1>

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

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.

### 使用

```python
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 训练脚本

```Python
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.  **泛化保留**：因为模型不需要在训练中改变它的“中文语言表征”（因为它不用预测中文），它在微调后，依然完美保留了预训练时学到的常识和逻辑能力，从而在推理时能听懂“大熊猫”和“等差数列”的区别。

<h1 align="center">Supra-Router-51M · Multi-Task Infrastructure Routing Model</h1>

![logo](https://cdn-uploads.huggingface.co/production/uploads/697f2832c2c5e4daa93cece7/lClDSdp9BkKyv_VGWSQbt.png)

<h2 align="center">About the Model</h2>

GGUF model [here](https://huggingface.co/SupraLabs/Supra-Router-51M-gguf)

**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:
```text
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.

```python
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.