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