docs: 重构 README 结构，全文档添加目录导航

- README 新增 Getting Started 端到端流程，整合快速开始与演示，去重精简 - 中文 README 同步英文版结构，预处理配置改用 seq 策略 - inference.md 补充 SSE 流式格式、错误响应、/stats 端点文档 - params.md 扩展为 CLI 参考，覆盖 server/generate/preprocess 参数表 - dataflow.md 拆分 tokenization/format detection/backend 子节，新增流程图 - architecture/training/inference/preprocessing 均添加目录导航 - 移除 README CI badge
docs: 修复文档中过时的字段、签名和缺失的类
2026-06-19 13:53:22 +08:00 · 2026-06-18 18:49:46 +08:00 · 2026-06-18 17:38:01 +08:00 · 2026-06-18 16:45:16 +08:00 · 2026-06-18 16:35:05 +08:00 · 2026-06-18 15:55:15 +08:00
18 changed files with 802 additions and 301 deletions
--- a/.github/workflows/update-badges.yml
+++ b/.github/workflows/update-badges.yml
@ -1,44 +0,0 @@
-name: Update Badges
-
-on:
-  push:
-    branches: [main]
-  schedule:
-    - cron: "0 0 * * *"
-  workflow_dispatch:
-
-permissions:
-  contents: write
-
-jobs:
-  update:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-
-      - name: Fetch repo stats
-        id: api
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          mkdir -p badges
-
-          REPO=$(gh repo view --json stargazerCount,forkCount,latestRelease --jq '.')
-
-          STARS=$(echo "$REPO" | jq -r '.stargazerCount')
-          FORKS=$(echo "$REPO" | jq -r '.forkCount')
-          RELEASE=$(echo "$REPO" | jq -r '.latestRelease.tagName // "N/A"')
-
-          echo '{"schemaVersion":1,"label":"release","message":"'"$RELEASE"'","color":"76bad9"}' > badges/release.json
-          echo '{"schemaVersion":1,"label":"stars","message":"'"$STARS"'","color":"76bad9"}' > badges/stars.json
-          echo '{"schemaVersion":1,"label":"forks","message":"'"$FORKS"'","color":"76bad9"}' > badges/forks.json
-
-      - name: Deploy to gh-pages
-        uses: peaceiris/actions-gh-pages@v4
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          publish_dir: badges
-          destination_dir: badges
-          commit_message: "Sync badges"
-          user_name: "github-actions[bot]"
-          user_email: "github-actions[bot]@users.noreply.github.com"
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@ -5,7 +5,7 @@ Thank you for your interest in contributing! This document provides step-by-step
 ## Quick Start

 ```bash
-git clone https://github.com/your-username/AstrAI.git
+git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e ".[dev]"     # install with dev dependencies (pytest, ruff)
 ```
--- a/README.md
+++ b/README.md
@ -9,9 +9,9 @@
 <div align="center">
  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/ViperEkura/AstrAI/gh-pages/badges/release.json" alt="release">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/ViperEkura/AstrAI/gh-pages/badges/stars.json" alt="stars">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/ViperEkura/AstrAI/gh-pages/badges/forks.json" alt="forks">
+  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
+  <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
 <br>

@ -28,7 +28,8 @@
 ## 📖 Table of Contents

 - [Features](#features)
- [Quick Start](#quick-start)
+- [Getting Started](#getting-started)
+- [Demo](#demo)
 - [Documentation](#documentation)
 - [Contributing](#contributing)
 - [Community](#community)
@ -49,33 +50,43 @@
 - 🤗 **HuggingFace-Style API**: AutoModel/AutoTokenizer APIs inspired by HuggingFace for easy model and tokenizer loading.
 - 🔌 **Dual API Compatibility**: Supports both OpenAI and Anthropic chat completion APIs out of the box.

-### Quick Start
+### Getting Started

-#### Installation
+End-to-end walkthrough in 5 steps:
+
+**1. Install**

 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e .
+# pip install -e ".[dev]"    # optional: dev dependencies (pytest, ruff)
 ```

-For development dependencies:
+**2. Download model**

 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # downloads 1B checkpoint to params/
 ```

-#### Download Pre-trained Model
+**3. Preprocess data**

-Download pre-trained model weights (1B bilingual checkpoint) to `params/`:
+Create `pretrain.json` (preprocessing config for `seq` strategy):
+
+```json
+{
+    "version": 1,
+    "input": {"sections": [{"field": "text", "action": "train"}]},
+    "preprocessing": {"max_seq_len": 2048},
+    "output": {"storage_format": "bin"}
+}
+```

 ```bash
-python scripts/demo/download.py
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```

-Or download manually from [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) into `params/`.
-
-#### Train a Model
+**4. Train**

 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
@ -102,15 +113,54 @@ nohup python scripts/tools/train.py \
    > out.log 2> err.log &
 ```

-Full reference at [Parameter Guide](assets/docs/params.md).
+**5. Serve & query**

-#### Generate Text
+```bash
+# Terminal 1: start server
+python scripts/tools/server.py --param_path ./params --device cuda
+
+# Terminal 2: query
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+### Demo
+
+Check out the demos in the `scripts/demo/` folder:
+
+```bash
+# Download model weights (required before running demos)
+python scripts/demo/download.py                      # model → params/
+
+# Interactive streaming chat (multi-turn, maintains history)
+python scripts/demo/stream_chat.py
+# Type your message after >>, type !exit to quit
+
+# Batch generation (5 hardcoded prompts, non-streaming)
+python scripts/demo/generate_batch.py
+
+# Single-prompt autoregressive streaming
+python scripts/demo/generate_ar.py
+```
+
+All generation demos use `temperature=0.8`, `top_p=0.95`, `top_k=50`, `max_tokens=2048` by default and require `params/` to contain model weights (run `download.py` first).
+
+Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
+
+---
+
+See [Documentation](#documentation) for full references beyond the examples above.
+
+#### Text Generation
+
+Batch generation from a JSONL file:

 ```bash
 python scripts/tools/generate.py \
-    --param_path /path/to/model \
-    --input_json_file /path/to/input.jsonl \
-    --output_json_file /path/to/output.jsonl
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
 ```

 #### Docker
@ -124,9 +174,6 @@ docker build -t astrai:latest .
 # Run with GPU support
 docker run --gpus all -it astrai:latest

-# Run with specific GPUs
-docker run --gpus '"device=0,1"' -it astrai:latest
-
 # Run inference server
 docker run --gpus all -p 8000:8000 astrai:latest \
  python -m scripts.tools.server --port 8000 --device cuda
@ -143,84 +190,37 @@ docker compose --profile cpu up -d

 > **Note**: `--gpus all` is required for CUDA support. Without it, `torch.cuda.is_available()` will return `False`.

-#### Start HTTP Server
+#### HTTP API Examples

-Start the inference server with OpenAI and Anthropic-compatible HTTP API:
+Additional request examples beyond the [Getting Started](#getting-started) flow:

 ```bash
-python -m scripts.tools.server --port 8000 --device cuda
-```
-
-Make requests:
-
-```bash
-# OpenAI-compatible
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Hello"}],
-    "max_tokens": 512
-  }'
-
 # OpenAI-compatible streaming
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Tell a story"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
+  -d '{"messages":[{"role":"user","content":"Tell a story"}],"stream":true,"max_tokens":500}'

 # Anthropic-compatible
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "You are a helpful assistant.",
-    "messages": [{"role": "user", "content": "Hello"}],
-    "max_tokens": 512
-  }'
+  -d '{"model":"astrai","system":"You are a helpful assistant.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'

 # Anthropic-compatible streaming with stop sequences
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "messages": [{"role": "user", "content": "Write a story"}],
-    "max_tokens": 500,
-    "stream": true,
-    "stop_sequences": ["The end"]
-  }'
+  -d '{"model":"astrai","messages":[{"role":"user","content":"Write a story"}],"max_tokens":500,"stream":true,"stop_sequences":["The end"]}'

 # Health check
 curl http://localhost:8000/health
 ```

-#### Demo
-
-Check out the demos in the `scripts/demo/` folder:
-
-```bash
-# Download model weights (required before running demos)
-python scripts/demo/download.py
-
-# Interactive streaming chat
-python scripts/demo/stream_chat.py
-
-# Batch generation
-python scripts/demo/generate_batch.py
-
-# Auto‑regressive generation
-python scripts/demo/generate_ar.py
-```
-
-Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
+See [Inference Guide](assets/docs/inference.md) for SSE streaming format, error codes, and stats endpoint.

 ### Documentation

 | Document | Description |
 |----------|-------------|
-| [Parameter Guide](./assets/docs/params.md) | Training & inference parameters |
+| [CLI Reference](./assets/docs/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) |
 | [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns |
 | [Training](./assets/docs/training.md) | Training loop, strategies & formulas |
 | [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API |
--- a/assets/docs/README-zh-CN.md
+++ b/assets/docs/README-zh-CN.md
@ -15,9 +15,9 @@
 <div align="center">
  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/ViperEkura/AstrAI/gh-pages/badges/release.json" alt="release">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/ViperEkura/AstrAI/gh-pages/badges/stars.json" alt="stars">
-  <img src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/ViperEkura/AstrAI/gh-pages/badges/forks.json" alt="forks">
+  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
+  <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>

 <br>
@ -34,7 +34,8 @@
 ## 📖 目录

 - [特性](#特性)
- [快速开始](#快速开始)
+- [快速上手](#快速上手)
+- [演示](#演示)
 - [文档](#文档)
 - [贡献](#贡献)
 - [社区](#社区)
@ -55,33 +56,43 @@
 - 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
 - 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。

-### 快速开始
+### 快速上手

-#### 安装
+端到端演示，只需 5 步：
+
+**1. 安装**

 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e .
+# pip install -e ".[dev]"    # 可选：开发依赖（pytest, ruff）
 ```

-安装开发依赖：
+**2. 下载模型**

 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # 下载 1B 检查点到 params/
 ```

-#### 下载预训练模型
+**3. 预处理数据**

-下载预训练模型权重（1B 双语检查点）到 `params/` 目录：
+创建 `pretrain.json`（`seq` 策略的预处理配置）：
+
+```json
+{
+    "version": 1,
+    "input": {"sections": [{"field": "text", "action": "train"}]},
+    "preprocessing": {"max_seq_len": 2048},
+    "output": {"storage_format": "bin"}
+}
+```

 ```bash
-python scripts/demo/download.py
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```

-或从 [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) 手动下载放入 `params/`。
-
-#### 训练模型
+**4. 训练**

 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
@ -108,15 +119,54 @@ nohup python scripts/tools/train.py \
    > out.log 2> err.log &
 ```

-完整参数列表见[参数说明](./params.md)。
+**5. 启动服务并调用**
+
+```bash
+# 终端 1：启动服务
+python scripts/tools/server.py --param_path ./params --device cuda
+
+# 终端 2：发起请求
+curl http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
+```
+
+### 演示
+
+查看 `scripts/demo/` 文件夹中的演示：
+
+```bash
+# 下载模型权重（运行演示前必需）
+python scripts/demo/download.py                      # model → params/
+
+# 交互式流式聊天（多轮对话，保持历史记录）
+python scripts/demo/stream_chat.py
+# 在 >> 后输入消息，输入 !exit 退出
+
+# 批量生成（5 条硬编码提示词，非流式）
+python scripts/demo/generate_batch.py
+
+# 单条提示词自回归流式生成
+python scripts/demo/generate_ar.py
+```
+
+所有生成演示默认使用 `temperature=0.8`、`top_p=0.95`、`top_k=50`、`max_tokens=2048`，需要 `params/` 目录包含模型权重（请先运行 `download.py`）。
+
+观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
+
+---
+
+更多选项请参考[文档](#文档)。

 #### 文本生成

+从 JSONL 文件批量生成：
+
 ```bash
 python scripts/tools/generate.py \
-    --param_path /path/to/model \
-    --input_json_file /path/to/input.jsonl \
-    --output_json_file /path/to/output.jsonl
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
 ```

 #### Docker
@ -130,9 +180,6 @@ docker build -t astrai:latest .
 # 启用 GPU 运行
 docker run --gpus all -it astrai:latest

-# 指定特定 GPU
-docker run --gpus '"device=0,1"' -it astrai:latest
-
 # 运行推理服务
 docker run --gpus all -p 8000:8000 astrai:latest \
  python -m scripts.tools.server --port 8000 --device cuda
@ -149,84 +196,37 @@ docker compose --profile cpu up -d

 > **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。

-#### 启动 HTTP 服务
+#### HTTP API 示例

-启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：
+除[快速上手](#快速上手)流程外，更多请求示例：

 ```bash
-python -m scripts.tools.server --port 8000 --device cuda
-```
-
-发起请求：
-
-```bash
-# OpenAI 兼容
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "你好"}],
-    "max_tokens": 512
-  }'
-
 # OpenAI 兼容流式
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "讲个故事"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
+  -d '{"messages":[{"role":"user","content":"讲个故事"}],"stream":true,"max_tokens":500}'

 # Anthropic 兼容
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "你是一个乐于助人的助手。",
-    "messages": [{"role": "user", "content": "你好"}],
-    "max_tokens": 512
-  }'
+  -d '{"model":"astrai","system":"你是一个乐于助人的助手。","messages":[{"role":"user","content":"你好"}],"max_tokens":512}'

 # Anthropic 兼容流式并设置停止序列
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "messages": [{"role": "user", "content": "写个故事"}],
-    "max_tokens": 500,
-    "stream": true,
-    "stop_sequences": ["结束"]
-  }'
+  -d '{"model":"astrai","messages":[{"role":"user","content":"写个故事"}],"max_tokens":500,"stream":true,"stop_sequences":["结束"]}'

 # 健康检查
 curl http://localhost:8000/health
 ```

-#### 演示
-
-查看 `scripts/demo/` 文件夹中的演示：
-
-```bash
-# 下载模型权重（运行演示前必需）
-python scripts/demo/download.py
-
-# 交互式流式聊天
-python scripts/demo/stream_chat.py
-
-# 批量生成
-python scripts/demo/generate_batch.py
-
-# 自回归生成
-python scripts/demo/generate_ar.py
-```
-
-观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
+SSE 流式格式、错误码和统计端点详见[推理文档](./inference.md)。

 ### 文档

 | 文档 | 说明 |
 |------|------|
-| [参数说明](./params.md) | 训练与推理参数配置 |
+| [CLI 参考](./params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
 | [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
 | [训练文档](./training.md) | 训练循环、策略与公式 |
 | [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
--- a/assets/docs/architecture.md
+++ b/assets/docs/architecture.md
@ -1,5 +1,12 @@
 # AstrAI Architecture

+## Contents
+
+- [Class Diagram](#class-diagram) — Full Mermaid class diagram across 10+ namespaces
+- [Module Overview](#module-overview) — Component inventory per module
+- [Design Patterns](#design-patterns) — 13 documented patterns with classes
+- [Core Relationships](#core-relationships) — 11 key inter-component relationships
+
 ## Class Diagram

 ```mermaid
@ -8,8 +15,8 @@ classDiagram
        class BaseConfig {
            +to_dict() Dict
            +from_dict(d) Self
-            +from_json(path) Self
-            +to_json(path)
+            +from_file(path) Self
+            +to_file(path)
        }

        class BaseModelConfig {
@ -61,31 +68,32 @@ classDiagram
        }

        class ConfigFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +load(raw) BaseConfig
        }

        class InputConfig {
-            +str type
-            +str messages_key
-            +str prompt_key
-            +str response_key
-            +str text_key
+            +Optional[List[Dict]] sections
+            +Optional[Dict[str, Dict]] sources
        }

        class ProcessingConfig {
            +int max_seq_len
            +int min_chars
            +int max_chars
-            +bool deduplicate
            +Optional[int] max_items
+            +str packing_strategy
+            +int max_packed_len
+            +str truncation_mode
        }

        class OutputConfig {
            +Optional[str] domain_key
            +str storage_format
            +int max_tokens_per_shard
+            +Dict[str, str] dtype
+            +str position_ids_mode
        }

        class PipelineConfig {
@ -190,13 +198,13 @@ classDiagram
        }

        class StoreFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +create(storage_type) Store
        }

        class DatasetFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +create(train_type, window_size, stride) BaseDataset
            +load(train_type, load_path, window_size, stride, storage_type) BaseDataset
@ -219,7 +227,7 @@ classDiagram
    namespace model {
        class AutoModel {
            +BaseModelConfig config
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +get_component_class(name) Type
            +from_pretrained(path, disable_random_init, strict) nn.Module
@ -395,24 +403,17 @@ classDiagram
    }

    namespace factory {
-        class Registry {
-            +Dict _entries
-            +register(name, component_cls, category, priority)
-            +get(name) Type
-            +list_names() List[str]
-        }
-
        class BaseFactory {
-            +Registry _registry
-            +register(name, category, priority) decorator
+            +Dict _entries
+            +register(name) decorator
            +create(name, *args, **kwargs) T
            +list_registered() list
        }

        class MaskBuilderFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
-            +create(input_type, config, tokenizer) BaseMaskBuilder
+            +create(name, *args, **kwargs) BaseMaskBuilder
        }
    }

@ -461,7 +462,7 @@ classDiagram
        }

        class StrategyFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +create(train_type, model, device, **kwargs) BaseStrategy
        }
@ -502,9 +503,9 @@ classDiagram
        }

        class SchedulerFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
-            +create(optimizer, schedule_type, **kwargs) BaseScheduler
+            +create(name, *args, **kwargs) BaseScheduler
        }

        class CosineScheduler {
@ -521,6 +522,13 @@ classDiagram
            +int t_mult
        }

+        class WSDScheduler {
+            +int warmup_steps
+            +int stable_steps
+            +int decay_steps
+            +float min_rate
+        }
+
        class TrainCallback {
            <<protocol>>
            +on_train_begin(context)
@ -581,7 +589,7 @@ classDiagram
        }

        class CallbackFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +create(name, **kwargs) TrainCallback
        }
@ -891,9 +899,9 @@ classDiagram
            +str yielded
        }

-        class app {
-            <<singleton>>
-            +FastAPI app
+        class get_app {
+            <<module>>
+            +get_app() FastAPI
        }
    }

@ -975,7 +983,7 @@ classDiagram
        }

        class ExecutorFactory {
-            +Registry _registry
+            +Dict _entries
            +register(name) decorator
            +create(parallel_mode, **kwargs) BaseExecutor
        }
@ -1018,6 +1026,7 @@ classDiagram
    BaseStrategy <|-- GRPOStrategy
    BaseScheduler <|-- CosineScheduler
    BaseScheduler <|-- SGDRScheduler
+    BaseScheduler <|-- WSDScheduler
    TrainCallback <|-- GradientClippingCallback
    TrainCallback <|-- GradientCheckpointingCallback
    TrainCallback <|-- CheckpointCallback
@ -1080,7 +1089,6 @@ classDiagram
    DecoderBlock *-- RMSNorm
    ChatCompletionRequest *-- ChatMessage
    MessagesRequest *-- AnthropicMessage
-    BaseFactory *-- Registry
    BaseExecutor *-- GradientState
    AccumOptimizer o-- GradientState
    AccumScheduler o-- GradientState
@ -1157,13 +1165,13 @@ classDiagram

 | Module | Components | Description |
 |--------|------------|-------------|
-| **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig, PipelineConfig, InputConfig, ProcessingConfig, OutputConfig | Configuration management (to_dict/from_dict, to_file/from_file, from_json/to_json) |
+| **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig, PipelineConfig, InputConfig, ProcessingConfig, OutputConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
 | **astrai.preprocessing** | BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, Pipeline, filter_by_length, PackingStrategy, PackingStrategyFactory, PositionIdStrategy, PositionIdStrategyFactory, StoreWriter, StoreWriterFactory | Declarative JSON-driven data preprocessing |
 | **astrai.dataset** | BaseDataset–GRPODataset, Store–MmapStore, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
 | **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
-| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–SGDRScheduler, SchedulerFactory, TrainCallback(Protocol)–ValidationCallback, CallbackFactory, Muon | Training workflow |
+| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–ValidationCallback, CallbackFactory, Muon | Training workflow |
 | **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–KvcacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, ChatMessage–MessagesRequest, app | Inference service |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation |
 | **astrai.factory** | Registry, BaseFactory[T] | Component registration |
@ -1174,7 +1182,7 @@ classDiagram
 | Pattern | Classes | Purpose |
 |---------|---------|---------|
 | **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory` | Decorator-based component creation |
-| **Registry** | `BaseFactory`, `Registry` | Component registration with category/priority |
+| **Registry** | `BaseFactory` | Component registration |
 | **Strategy** | `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy` | Training strategy switching |
 | **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations |
 | **Strategy (API)** | `ResponseBuilder`, `OpenAIResponseBuilder`, `AnthropicResponseBuilder` | HTTP API handler with format hooks |
@ -1197,7 +1205,7 @@ classDiagram
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore) loads data with explicit `_length` and multi-segment `_data`
 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt`
-9. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`
+9. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`/`WSDScheduler`
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers

--- a/assets/docs/dataflow.md
+++ b/assets/docs/dataflow.md
@ -1,17 +1,58 @@
 # Data Flow

-This document describes the data pipeline: from raw text to model input tensors.
+This document describes the data pipeline: from raw text to model input tensors. For creating preprocessing configs, see [Preprocessing Guide](preprocessing.md).
+
+## Contents
+
+- [Overview](#overview)
+- [Data Preparation](#data-preparation) — tokenization, format detection, backends
+- [Data Keys by Training Type](#data-keys-by-training-type)
+- [Dataset Architecture](#dataset-architecture)
+- [Sampler](#sampler)
+- [DataLoader](#dataloader)

 ## Overview

 ```
-Raw Text → AutoTokenizer → Token IDs → .h5/.bin → Store.load() → Store.fetch() → Dataset → Sampler → DataLoader → Training/Inference
+JSONL Lines → Pipeline (mask builder) → Tokenized Tensors
+                                              ↓
+                                      .h5 or .bin storage
+                                              ↓
+                                      Store.load()
+                                              ↓
+                                      Store.fetch(begin, end, keys)
+                                              ↓
+                                      BaseDataset.__getitem__(idx)
+                                              ↓
+                                      Sampler → DataLoader → Training / Inference
 ```

 ## Data Preparation

 Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.

+### Tokenization
+
+The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](preprocessing.md)), and produces flat token sequences:
+
+```python
+# Per JSONL line: messages → chat template → token IDs + loss mask
+tokens = tokenizer.encode(rendered_text)        # List[int]
+loss_mask = [0, 0, 0, 1, 1, 1, 1, 1, 1]        # 0=masked, 1=train
+# Stored as flat tensors, packed with other lines by packing strategy
+```
+
+The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
+
+### Format Detection
+
+`detect_format(load_path)` inspects the directory:
+
+- If `*.h5` files exist → `"h5"` (HDF5 backend)
+- If `*.bin` + `meta.json` files exist → `"bin"` (memory-mapped backend)
+
+### Store Backends
+
 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:

 ```
@ -19,7 +60,11 @@ StoreFactory.create("h5")  → H5Store
 StoreFactory.create("bin") → MmapStore
 ```

-H5 backend supports shared memory via `.share_memory_()`. Bin (mmap) uses OS page-cache sharing natively.
+**H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory).
+
+**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
+
+Both backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing).

 ## Data Keys by Training Type

@ -61,4 +106,4 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ

 Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.

-> Document Update Time: 2026-05-30
+> Document Update Time: 2026-06-19
--- a/assets/docs/inference.md
+++ b/assets/docs/inference.md
@ -1,5 +1,16 @@
 # Inference

+## Contents
+
+- [KV Cache](#kv-cache)
+- [KVCache System](#kvcache-system)
+- [Continuous Batching](#continuous-batching)
+- [Sampling](#sampling-strategy-pattern)
+- [Protocol Handlers](#protocol-handlers-strategy-pattern)
+- [Engine & GenerateResult](#engine--generateresult)
+- [HTTP API](#http-api) — endpoints, SSE, errors, stats
+- [Engine API](#engine-api)
+
 ## KV Cache

 At decode time, only the last query token matters. All previous K/V are cached to avoid recomputation:
@ -133,6 +144,92 @@ Supports `stop_sequences` and streaming via `event: content_block_delta`.
 | `max_tokens` | Optional[int] | None | Max generation length |
 | `stream` | bool | False | Stream output |

+### SSE Streaming Format
+
+**OpenAI** (`/v1/chat/completions`, `stream=true`):
+
+```
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
+       "choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk",...,
+       "choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk",...,
+       "choices":[{"index":0,"delta":{},"finish_reason":"stop"}],
+       "usage":{"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}}
+
+data: [DONE]
+```
+
+**Anthropic** (`/v1/messages`, `stream=true`):
+
+```
+event: message_start
+data: {"type":"message_start","message":{"id":"msg_...","model":"astrai","role":"assistant",
+       "content":[],"stop_reason":null,...}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{...}}
+
+event: message_stop
+data: {"type":"message_stop"}
+```
+
+### Error Responses
+
+All endpoints use standard HTTP status codes:
+
+| Status | Meaning |
+|--------|---------|
+| 200 | Success |
+| 400 | Invalid request (bad JSON, missing fields, validation error) |
+| 405 | Method not allowed |
+| 422 | Unprocessable entity (Pydantic validation) |
+| 500 | Internal server error (model crash, OOM, scheduler failure) |
+| 503 | Service unavailable (model not loaded, engine not ready) |
+
+Error response body:
+
+```json
+{
+    "error": {
+        "message": "Invalid request: max_tokens must be > 0",
+        "type": "invalid_request_error",
+        "code": 400
+    }
+}
+```
+
+### Stats Endpoint
+
+```
+GET /stats
+```
+
+Response:
+
+```json
+{
+    "active_requests": 3,
+    "waiting_requests": 2,
+    "total_requests": 128,
+    "cache_usage": 0.45,
+    "tokens_generated": 10240
+}
+```
+
+`cache_usage` is the fraction of KV cache pages currently in use (0.0–1.0).
+
 ## Engine API

 ```python
@ -149,4 +246,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
    print(token)
 ```

-> Document Update Time: 2026-05-30
+> Document Update Time: 2026-06-19
--- a/assets/docs/params.md
+++ b/assets/docs/params.md
@ -1,4 +1,11 @@
-# Parameter Documentation
+# CLI Parameter Reference
+
+## Contents
+
+- [Training Parameters](#training-parameters)
+- [Inference Server](#inference-server-serverpy)
+- [Generate](#generate-generatepy)
+- [Preprocess](#preprocess-preprocesspy)

 ## Training Parameters

@ -86,11 +93,12 @@
 | Parameter | Description | Default | Used by |
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
-| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.05 | `seq`, `sft` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` |
 | `--group_size` | GRPO group size | 4 | `grpo` |
 | `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
 | `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
 | `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
+| `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |

 ### Usage Example

@ -121,4 +129,64 @@ nohup python scripts/tools/train.py \

 ---

-> Document Update Time: 2026-05-24
+## Inference Server (`server.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--host` | str | `0.0.0.0` | Host address |
+| `--port` | int | `8000` | Port number |
+| `--param_path` | path | `project_root/params` | Path to model parameters |
+| `--device` | str | `cuda` | Device to load model on |
+| `--dtype` | str | `bfloat16` | Model weights dtype (`bfloat16`, `float16`, `float32`) |
+| `--max_batch_size` | int | `16` | Maximum batch size for continuous batching |
+| `--reload` | flag | `False` | Enable auto-reload for development |
+
+Usage:
+```bash
+python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloat16
+```
+
+See [Inference Guide](inference.md) for HTTP API documentation.
+
+## Generate (`generate.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--param_path` | str | required | Path to the model directory |
+| `--input_json_file` | str | required | Path to the input JSONL file |
+| `--output_json_file` | str | required | Path to the output JSONL file |
+| `--question_key` | str | `question` | Key for the question in input JSON |
+| `--response_key` | str | `response` | Key for the response in output JSON |
+| `--temperature` | float | `0.60` | Sampling temperature |
+| `--top_k` | int | `30` | Top-k filtering |
+| `--top_p` | float | `0.95` | Nucleus sampling threshold |
+| `--batch_size` | int | `1` | Batch size for generation |
+| `--max_tokens` | int | `2048` | Maximum tokens to generate |
+
+Usage:
+```bash
+python scripts/tools/generate.py \
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
+```
+
+## Preprocess (`preprocess.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) |
+| `--output_dir`, `-o` | path | required | Output directory for processed data |
+| `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) |
+| `--num_workers` | int | `4` | Number of parallel workers |
+
+Usage:
+```bash
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c sft.json
+```
+
+See [Preprocessing Guide](preprocessing.md) for config file format and examples.
+
+---
+
+> Document Update Time: 2026-06-19
--- a/assets/docs/preprocessing.md
+++ b/assets/docs/preprocessing.md
@ -2,6 +2,17 @@

 Declarative JSON-driven data preprocessing. One `SectionedMaskBuilder` handles all formats via `input.sections` (single-output) or `input.sources` (multi-output).

+## Contents
+
+- [Philosophy](#philosophy)
+- [Config Structure](#config-structure)
+- [Quick Start](#quick-start) — SFT Chat, SFT Instruction, Pretrain, DPO, GRPO examples
+- [Configuration Reference](#configuration-reference) — all fields
+- [Mask Algorithm](#mask-algorithm)
+- [Output Layout](#output-layout)
+- [CLI](#cli)
+- [Python API](#python-api)
+
 ## Philosophy

 | Component | Responsibility |
--- a/assets/docs/training.md
+++ b/assets/docs/training.md
@ -1,5 +1,18 @@
 # Training

+## Contents
+
+- [Autoregression](#autoregression)
+- [Causal Mask](#causal-mask)
+- [Rotary Position Embedding (RoPE)](#rotary-position-embedding-rope)
+- [Training Loop](#training-loop)
+- [Strategies](#strategies) — SEQ, SFT, DPO, GRPO
+- [LR Schedulers](#lr-schedulers)
+- [Gradient Checkpointing](#gradient-checkpointing)
+- [Checkpoint](#checkpoint)
+- [TrainContextBuilder](#traincontextbuilder-builder-pattern)
+- [Training CLI](#training-cli)
+
 ### Autoregression

 Given a token sequence, the model predicts the probability of the next token. Each generated token is appended to the input and fed back, repeating until an end-of-sequence token or max length.
@ -127,8 +140,9 @@ Keys: `prompts`, `responses`, `masks`, `rewards`.
 |------|-------|-------------|
 | Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
 | SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
+| WSD | `WSDScheduler` | Warmup-Stable-Decay with sqrt cooldown |

-Created by `SchedulerFactory.create(optimizer, schedule_type, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`. Omit to use no scheduler.
+Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. Omit to use no scheduler.

 ## Gradient Checkpointing

--- a/astrai/init.py
+++ b/astrai/init.py
@ -3,32 +3,98 @@ __author__ = "ViperEkura"

 from astrai.config import (
    AutoRegressiveLMConfig,
+    BaseModelConfig,
+    ConfigFactory,
    EncoderConfig,
+    PipelineConfig,
    TrainConfig,
 )
-from astrai.dataset import DatasetFactory
+from astrai.dataset import (
+    BaseDataset,
+    DatasetFactory,
+    ResumableDistributedSampler,
+    Store,
+    StoreFactory,
+)
 from astrai.factory import BaseFactory
 from astrai.inference import (
    GenerationRequest,
    InferenceEngine,
+    ProtocolHandler,
+    SamplingPipeline,
+    get_app,
+    run_server,
+    sample,
+)
+from astrai.model import (
+    AutoModel,
+    AutoRegressiveLM,
+    EmbeddingEncoder,
+    LoRAConfig,
+    inject_lora,
+)
+from astrai.parallel import (
+    ExecutorFactory,
+    get_rank,
+    get_world_size,
+    only_on_rank,
+    spawn_parallel_fn,
+)
+from astrai.preprocessing import Pipeline, filter_by_length
+from astrai.serialization import Checkpoint
+from astrai.tokenize import AutoTokenizer, ChatTemplate
+from astrai.trainer import (
+    BaseScheduler,
+    BaseStrategy,
+    CallbackFactory,
+    Muon,
+    SchedulerFactory,
+    StrategyFactory,
+    TrainCallback,
+    Trainer,
 )
-from astrai.model import AutoModel, AutoRegressiveLM
-from astrai.tokenize import AutoTokenizer
-from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer

 __all__ = [
    "AutoRegressiveLM",
    "AutoRegressiveLMConfig",
-    "EncoderConfig",
-    "TrainConfig",
-    "DatasetFactory",
+    "AutoModel",
    "AutoTokenizer",
+    "BaseDataset",
+    "BaseFactory",
+    "BaseModelConfig",
+    "BaseScheduler",
+    "BaseStrategy",
+    "CallbackFactory",
+    "ChatTemplate",
+    "Checkpoint",
+    "ConfigFactory",
+    "DatasetFactory",
+    "EmbeddingEncoder",
+    "EncoderConfig",
+    "ExecutorFactory",
    "GenerationRequest",
    "InferenceEngine",
-    "Trainer",
-    "CallbackFactory",
-    "StrategyFactory",
+    "LoRAConfig",
+    "Muon",
+    "Pipeline",
+    "PipelineConfig",
+    "ProtocolHandler",
+    "ResumableDistributedSampler",
+    "SamplingPipeline",
    "SchedulerFactory",
-    "BaseFactory",
-    "AutoModel",
+    "Store",
+    "StoreFactory",
+    "StrategyFactory",
+    "TrainCallback",
+    "TrainConfig",
+    "Trainer",
+    "filter_by_length",
+    "get_app",
+    "get_rank",
+    "get_world_size",
+    "inject_lora",
+    "only_on_rank",
+    "run_server",
+    "sample",
+    "spawn_parallel_fn",
 ]
--- a/astrai/inference/core/scheduler.py
+++ b/astrai/inference/core/scheduler.py
@ -70,8 +70,8 @@ class InferenceScheduler:
            dtype=self.dtype,
        )

-        self._running = False
-        self._fatal_error: Optional[Exception] = None
+        self._stop_event = threading.Event()
+        self._loop_thread: Optional[threading.Thread] = None

    def add_task(self, prompt: str, **kwargs) -> str:
        return self._task_mgr.add_task(prompt, **kwargs)
@ -86,7 +86,7 @@ class InferenceScheduler:
    def _run_generation_loop(self):
        stop_ids = self._task_mgr.tokenizer.stop_ids
        try:
-            while self._running:
+            while not self._stop_event.is_set():
                finished = self._task_mgr.remove_finished_tasks(stop_ids)
                for task in finished:
                    self._page_cache.task_free(task.task_id)
@ -176,8 +176,7 @@ class InferenceScheduler:
                                    t.stream_callback(STOP)

        except Exception as e:
-            self._fatal_error = e
-            self._running = False
+            self._stop_event.set()
            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
            for task in self._task_mgr.get_active_tasks():
                if task.stream_callback:
@ -189,17 +188,19 @@ class InferenceScheduler:
            self._task_mgr.clear_queues()

    def start(self):
-        if not self._running:
-            self._running = True
-            t = threading.Thread(target=self._run_generation_loop, daemon=True)
-            t.start()
-            self._loop_thread = t
+        if self._loop_thread is not None and self._loop_thread.is_alive():
+            return
+        self._stop_event.clear()
+        t = threading.Thread(target=self._run_generation_loop, daemon=True)
+        t.start()
+        self._loop_thread = t

    def stop(self):
-        self._running = False
+        self._stop_event.set()
        self._task_mgr.wake()
-        if hasattr(self, "_loop_thread"):
+        if self._loop_thread is not None:
            self._loop_thread.join(timeout=2.0)
+        self._loop_thread = None
        for task in self._task_mgr.get_active_tasks():
            if task.stream_callback:
                task.stream_callback(STOP)
--- a/astrai/preprocessing/packing.py
+++ b/astrai/preprocessing/packing.py
@ -6,8 +6,7 @@ pipeline later flattens the result into contiguous tensors.
 """

 from abc import ABC, abstractmethod
-from collections import defaultdict
-from typing import Dict, List, Tuple
+from typing import Dict, List

 from astrai.factory import BaseFactory

@ -53,6 +52,15 @@ class SimplePacking(PackingStrategy):

@PackingStrategyFactory.register("bfd")
 class BFDPacking(PackingStrategy):
+    """Best-Fit Decreasing bin packing.
+
+    Assigns sequences to bins using a best-fit heuristic (sorted by
+    decreasing length) and concatenates sequences within each bin into
+    a single packed sequence.  Packed sequences are truncated to
+    *max_packed_len* so that each packed bin fits within one context
+    window during training.
+    """
+
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
@ -62,24 +70,40 @@ class BFDPacking(PackingStrategy):
        sequences = keys.get("sequence", [])
        if not sequences:
            return keys
-        plan = self._plan(sequences, max_packed_len)
-        reordered: dict = defaultdict(list)
-        for orig_idx, _ in plan:
-            for k, vals in keys.items():
-                reordered[k].append(
-                    _truncate(vals[orig_idx], max_packed_len, truncation_mode)
+        bins = self._plan(sequences, max_packed_len, truncation_mode)
+
+        packed: Dict[str, List[List[int]]] = {}
+        for k, vals in keys.items():
+            packed[k] = [
+                _truncate(
+                    self._concat_bin(vals, bin_indices),
+                    max_packed_len,
+                    truncation_mode,
                )
-        return dict(reordered)
+                for bin_indices in bins
+            ]
+        return packed

    @staticmethod
-    def _plan(sequences: List[List[int]], max_packed_len: int) -> List[Tuple[int, int]]:
+    def _concat_bin(vals: List[List[int]], indices: List[int]) -> List[int]:
+        result: List[int] = []
+        for i in indices:
+            result.extend(vals[i])
+        return result
+
+    @staticmethod
+    def _plan(
+        sequences: List[List[int]], max_packed_len: int, truncation_mode: str
+    ) -> List[List[int]]:
        n = len(sequences)
        order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
        bins: List[List[int]] = []
        bin_lengths: List[int] = []

        for orig_idx in order:
-            seq_len = min(len(sequences[orig_idx]), max_packed_len)
+            seq_len = len(
+                _truncate(sequences[orig_idx], max_packed_len, truncation_mode)
+            )
            best_bin = None
            best_remain = max_packed_len + 1
            for i, bl in enumerate(bin_lengths):
@ -94,8 +118,4 @@ class BFDPacking(PackingStrategy):
                bins.append([orig_idx])
                bin_lengths.append(seq_len)

-        plan: List[Tuple[int, int]] = []
-        for bin_indices in bins:
-            for orig_idx in bin_indices:
-                plan.append((orig_idx, min(len(sequences[orig_idx]), max_packed_len)))
-        return plan
+        return bins
--- a/astrai/preprocessing/pipeline.py
+++ b/astrai/preprocessing/pipeline.py
@ -7,6 +7,7 @@ dispatched by configuration keys.
 """

 import json
+import logging
 import os
 from collections import defaultdict
 from itertools import chain
@ -22,6 +23,8 @@ from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.preprocessing.writer import StoreWriterFactory
 from astrai.tokenize import AutoTokenizer

+logger = logging.getLogger(__name__)
+
 _STR_TO_DTYPE: dict[str, torch.dtype] = {
    "bool": torch.bool,
    "uint8": torch.uint8,
@ -88,7 +91,13 @@ class Pipeline:
            if pp.max_items and count >= pp.max_items:
                break

-            result = self.transform(item)
+            try:
+                result = self.transform(item)
+            except Exception:
+                logger.warning(
+                    "Failed to process item #%d, skipping", count + 1, exc_info=True
+                )
+                continue
            if result is None:
                continue

@ -105,7 +114,7 @@ class Pipeline:
                continue

            bucket = domains[domain]
-            self._align_bucket(bucket, result, ids, is_multi)
+            self._align_bucket(bucket, result, ids)
            for key, val in result.items():
                bucket[key].append(val)

@ -130,16 +139,12 @@ class Pipeline:
        return []

    @staticmethod
-    def _align_bucket(bucket: dict, result: dict, ids: list, is_multi: bool):
+    def _align_bucket(bucket: dict, result: dict, ids: list):
        """Pad previously-accumulated keys that are missing from *result*."""
        for key in list(bucket.keys()):
            if key in result:
                continue
-            if is_multi:
-                pad = bucket[key][-1] if bucket[key] else [1] * len(ids)
-                bucket[key].append(pad)
-            else:
-                bucket[key].append([1] * len(ids))
+            bucket[key].append([1] * len(ids))

    def _iter_items(self):
        for path in self.paths:
--- a/astrai/preprocessing/writer.py
+++ b/astrai/preprocessing/writer.py
@ -6,7 +6,9 @@ List[Tensor]}`` dict and delegates the write to the writer selected
 by ``output.storage_format``.
 """

+import logging
 import os
+import shutil
 from abc import ABC, abstractmethod
 from typing import Dict, List

@ -15,6 +17,8 @@ import torch
 from astrai.dataset.storage import save_bin, save_h5
 from astrai.factory import BaseFactory

+logger = logging.getLogger(__name__)
+

 class StoreWriter(ABC):
    """Write pre-tokenized tensors to disk in a format-specific way."""
@ -37,11 +41,35 @@ class StoreWriterFactory(BaseFactory["StoreWriter"]):
 class BinWriter(StoreWriter):
    def save(self, output_dir, domain, shard_idx, tensors):
        shard_path = os.path.join(output_dir, domain, f"shard_{shard_idx:04d}")
-        save_bin(shard_path, tensors)
+        try:
+            save_bin(shard_path, tensors)
+        except Exception:
+            if os.path.exists(shard_path):
+                shutil.rmtree(shard_path, ignore_errors=True)
+            logger.error(
+                "Failed to write shard %s/%s_%04d, cleaned up partial output",
+                domain,
+                "shard",
+                shard_idx,
+                exc_info=True,
+            )
+            raise


@StoreWriterFactory.register("h5")
 class H5Writer(StoreWriter):
    def save(self, output_dir, domain, shard_idx, tensors):
        chunk_dir = os.path.join(output_dir, domain)
-        save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
+        file_path = os.path.join(chunk_dir, f"data_{shard_idx:04d}.h5")
+        try:
+            save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
+        except Exception:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            logger.error(
+                "Failed to write shard %s/data_%04d.h5, cleaned up partial output",
+                domain,
+                shard_idx,
+                exc_info=True,
+            )
+            raise
--- a/astrai/trainer/schedule.py
+++ b/astrai/trainer/schedule.py
@ -164,3 +164,66 @@ class SGDRScheduler(BaseScheduler):
        self.min_rate = state_dict.pop("min_rate")
        self.t_mult = state_dict.pop("t_mult")
        super().load_state_dict(state_dict)
+
+
+@SchedulerFactory.register("wsd")
+class WSDScheduler(BaseScheduler):
+    """WSD (Warmup-Stable-Decay) scheduler with sqrt cooldown.
+
+    warmup_steps:  linear warmup from min_rate to 1.0
+    stable_steps:   constant at base_lr
+    decay_steps:    sqrt decay from base_lr to min_rate
+    min_rate:       minimum lr as fraction of base_lr (default 0.0)
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int,
+        stable_steps: int,
+        decay_steps: int,
+        min_rate: float = 0.0,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.stable_steps = stable_steps
+        self.decay_steps = decay_steps
+        self.min_rate = min_rate
+        self.total_steps = warmup_steps + stable_steps + decay_steps
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> List[float]:
+        if self.last_epoch < self.warmup_steps:
+            factor = self.last_epoch / max(self.warmup_steps, 1)
+            return [base_lr * factor for base_lr in self.base_lrs]
+
+        offset = self.last_epoch - self.warmup_steps
+
+        if offset < self.stable_steps:
+            return list(self.base_lrs)
+
+        decay_ratio = (offset - self.stable_steps) / max(self.decay_steps, 1)
+        decay_ratio = min(decay_ratio, 1.0)
+        factor = (1.0 - self.min_rate) * (1.0 - decay_ratio) ** 2 + self.min_rate
+        return [base_lr * factor for base_lr in self.base_lrs]
+
+    def state_dict(self):
+        state = super().state_dict()
+        state.update(
+            {
+                "warmup_steps": self.warmup_steps,
+                "stable_steps": self.stable_steps,
+                "decay_steps": self.decay_steps,
+                "min_rate": self.min_rate,
+                "total_steps": self.total_steps,
+            }
+        )
+        return state
+
+    def load_state_dict(self, state_dict):
+        self.warmup_steps = state_dict.pop("warmup_steps")
+        self.stable_steps = state_dict.pop("stable_steps")
+        self.decay_steps = state_dict.pop("decay_steps")
+        self.min_rate = state_dict.pop("min_rate")
+        self.total_steps = state_dict.pop("total_steps")
+        super().load_state_dict(state_dict)
--- a/scripts/eval/evaluate_ifd.py
+++ b/scripts/eval/evaluate_ifd.py
@ -15,6 +15,13 @@ Usage::
    python scripts/eval/ifd.py --param_path ./params \
        --input data.jsonl --output data_with_ifd.jsonl \
        --instr_key instruction --resp_key response
+
+Disable chat template::
+
+    python scripts/eval/ifd.py --param_path ./params \
+        --input data.jsonl --output data_with_ifd.jsonl \
+        --instr_key instruction --resp_key response \
+        --no_chat_template
 """

 import argparse
@ -35,7 +42,16 @@ def compute_ifd(
    response: str,
    device: str,
    max_len: int = 2048,
+    use_chat_template: bool = False,
 ) -> dict:
+    if use_chat_template:
+        return _compute_ifd_with_template(
+            model, tokenizer, instruction, response, device, max_len
+        )
+    return _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len)
+
+
+def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -> dict:
    instr_ids = tokenizer.encode(instruction)
    resp_ids = tokenizer.encode(response)

@ -47,7 +63,6 @@ def compute_ifd(
            "error": "empty response",
        }

-    # Truncate instruction if total length exceeds max_len
    qa_len = len(instr_ids) + len(resp_ids)
    if qa_len > max_len:
        overflow = qa_len - max_len
@ -56,24 +71,22 @@ def compute_ifd(
    instr_len = len(instr_ids)
    resp_len = len(resp_ids)

-    # Conditional: instruction + response
    qa_ids = instr_ids + resp_ids
    qa_tensor = torch.tensor([qa_ids], device=device, dtype=torch.long)

    with torch.inference_mode():
-        logits_qa = model(qa_tensor)["logits"][0]  # [qa_len, vocab]
+        logits_qa = model(qa_tensor)["logits"][0]

-    resp_logits = logits_qa[instr_len - 1 : -1]  # predict response tokens
+    resp_logits = logits_qa[instr_len - 1 : -1]
    resp_targets = torch.tensor(resp_ids, device=device, dtype=torch.long)
    L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()

-    # Unconditional: response alone
    resp_tensor = torch.tensor([resp_ids], device=device, dtype=torch.long)

    with torch.inference_mode():
-        logits_resp = model(resp_tensor)["logits"][0]  # [resp_len, vocab]
+        logits_resp = model(resp_tensor)["logits"][0]

-    unp_logits = logits_resp[:-1]  # causal shift
+    unp_logits = logits_resp[:-1]
    unp_targets = resp_tensor[0, 1:]
    L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()

@ -89,6 +102,83 @@ def compute_ifd(
    }


+def _compute_ifd_with_template(
+    model, tokenizer, instruction, response, device, max_len
+) -> dict:
+    instr_prefix = tokenizer.apply_chat_template(
+        [{"role": "user", "content": instruction}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+    full_text = tokenizer.apply_chat_template(
+        [
+            {"role": "user", "content": instruction},
+            {"role": "assistant", "content": response},
+        ],
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+
+    full_ids = tokenizer.encode(full_text)
+    prefix_ids = tokenizer.encode(instr_prefix)
+    resp_ids = tokenizer.encode(response)
+
+    if not resp_ids:
+        return {
+            "L_cond": None,
+            "L_uncond": None,
+            "ifd": None,
+            "error": "empty response",
+        }
+
+    if len(full_ids) > max_len:
+        overflow = len(full_ids) - max_len
+        full_ids = full_ids[overflow:]
+        prefix_len = len(prefix_ids) - overflow
+        prefix_len = max(0, prefix_len)
+    else:
+        prefix_len = len(prefix_ids)
+
+    cond_tensor = torch.tensor([full_ids], device=device, dtype=torch.long)
+
+    with torch.inference_mode():
+        logits_qa = model(cond_tensor)["logits"][0]
+
+    resp_start = prefix_len - 1
+    resp_end = len(full_ids) - 1
+    if resp_end <= resp_start:
+        return {
+            "L_cond": None,
+            "L_uncond": None,
+            "ifd": None,
+            "error": "response truncated entirely",
+        }
+
+    resp_logits = logits_qa[resp_start:resp_end]
+    resp_targets = torch.tensor(full_ids[prefix_len:], device=device, dtype=torch.long)
+    L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
+
+    resp_tensor = torch.tensor([resp_ids], device=device, dtype=torch.long)
+
+    with torch.inference_mode():
+        logits_resp = model(resp_tensor)["logits"][0]
+
+    unp_logits = logits_resp[:-1]
+    unp_targets = resp_tensor[0, 1:]
+    L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
+
+    ifd = L_cond / L_uncond if L_uncond > 0 else None
+
+    return {
+        "L_cond": round(L_cond, 6),
+        "L_uncond": round(L_uncond, 6),
+        "ifd": round(ifd, 6) if ifd is not None else None,
+        "instr_len": prefix_len,
+        "resp_len": len(resp_ids),
+        "error": None,
+    }
+
+
 def process_file(
    param_path: str,
    input_file: str,
@ -96,6 +186,7 @@ def process_file(
    instr_key: str,
    resp_key: str,
    max_len: int,
+    use_chat_template: bool = False,
 ):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 if device == "cuda" else torch.float32
@ -105,6 +196,12 @@ def process_file(
    model.to(device=device, dtype=dtype)
    model.eval()

+    if use_chat_template and tokenizer._chat_template is None:
+        raise RuntimeError(
+            "--use_chat_template specified but tokenizer has no chat template. "
+            "Add a chat_template to tokenizer_config.json or omit the flag."
+        )
+
    with open(input_file, "r", encoding="utf-8") as f:
        data = [json.loads(line) for line in f if line.strip()]

@ -116,7 +213,13 @@ def process_file(
            instruction = item[instr_key]
            response = item[resp_key]
            scores = compute_ifd(
-                model, tokenizer, instruction, response, device, max_len
+                model,
+                tokenizer,
+                instruction,
+                response,
+                device,
+                max_len,
+                use_chat_template=use_chat_template,
            )
            ifd_values.append(scores["ifd"])
            results.append({**item, "ifd": scores["ifd"], "ifd_detail": scores})
@ -167,6 +270,12 @@ def main():
        default=2048,
        help="Max token length (instruction truncated to fit)",
    )
+    parser.add_argument(
+        "--no_chat_template",
+        action="store_true",
+        default=False,
+        help="Disable chat template, use raw text concatenation",
+    )
    args = parser.parse_args()

    process_file(
@ -176,6 +285,7 @@ def main():
        args.instr_key,
        args.resp_key,
        args.max_len,
+        use_chat_template=not args.no_chat_template,
    )


--- a/scripts/eval/evaluate_ifeval.py
+++ b/scripts/eval/evaluate_ifeval.py
@ -343,14 +343,20 @@ def verify_response(response: str, instruction_id: str, kwargs: dict) -> Optiona

 def generate_one(
    engine: InferenceEngine,
+    tokenizer: AutoTokenizer,
    prompt: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
 ) -> str:
+    formatted = tokenizer.apply_chat_template(
+        [{"role": "user", "content": prompt}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
    output = engine.generate(
-        prompt=prompt,
+        prompt=formatted,
        stream=False,
        max_tokens=max_tokens,
        temperature=temperature,
@ -364,6 +370,7 @@ def generate_one(

 def evaluate(
    engine: InferenceEngine,
+    tokenizer: AutoTokenizer,
    problems: List[dict],
    max_tokens: int,
    temperature: float,
@ -385,7 +392,7 @@ def evaluate(
        samples = []
        for _ in range(num_samples):
            response = generate_one(
-                engine, prompt, max_tokens, temperature, top_p, top_k
+                engine, tokenizer, prompt, max_tokens, temperature, top_p, top_k
            )
            samples.append(response)

@ -536,6 +543,7 @@ def main():
    model = AutoModel.from_pretrained(args.param_path)
    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
    model.to(device="cuda", dtype=torch.bfloat16)
+    model.eval()

    engine = InferenceEngine(
        model=model,
@ -545,6 +553,7 @@ def main():

    results = evaluate(
        engine=engine,
+        tokenizer=tokenizer,
        problems=problems,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
Author	SHA1	Message	Date
ViperEkura	8ab7564d02	docs: 重构 README 结构，全文档添加目录导航 - README 新增 Getting Started 端到端流程，整合快速开始与演示，去重精简 - 中文 README 同步英文版结构，预处理配置改用 seq 策略 - inference.md 补充 SSE 流式格式、错误响应、/stats 端点文档 - params.md 扩展为 CLI 参考，覆盖 server/generate/preprocess 参数表 - dataflow.md 拆分 tokenization/format detection/backend 子节，新增流程图 - architecture/training/inference/preprocessing 均添加目录导航 - 移除 README CI badge	2026-06-19 13:53:22 +08:00
ViperEkura	d096b6e29e	docs: 修复文档中过时的字段、签名和缺失的类 - BaseConfig 的 from_json/to_json → from_file/to_file - InputConfig/ProcessingConfig/OutputConfig 字段对齐源码 - 移除不存在的 Registry 类，register() 去 category/priority - SchedulerFactory.create 参数顺序修正 - 架构图/训练/参数文档补全 WSDScheduler - CONTRIBUTING.md 克隆地址占位符修正 - params.md label_smoothing 默认值修正，补全 neftune_alpha - app 类更正为 get_app 函数	2026-06-18 18:49:46 +08:00
ViperEkura	d88a41f8f1	fix: 修复预处理流水线 4 个致命问题 - pipeline: 单条数据异常不再崩溃整条流水线, 改 log warning 后跳过 - pipeline: _align_bucket 统一用 len(ids) 填充, 修复多输出模式下长度错配 - writer: BinWriter/H5Writer 写入失败自动清理残留文件并记录详细错误 - packing: BFDPacking 真正将序列打包进 bin 而非仅重排, 减少碎片	2026-06-18 17:38:01 +08:00
ViperEkura	376e9eba80	feat: IFEval 使用 chat template 格式化 prompt，添加 model.eval() - generate_one 用 tokenizer.apply_chat_template 包 user 消息 - 新增 model.eval() 关闭 dropout，确保确定性输出	2026-06-18 16:45:16 +08:00
ViperEkura	a62c2e11a2	feat: IFD 默认使用 chat template，支持裸文本模式 - 新增 _compute_ifd_with_template，用 tokenizer chat template 格式化后计算 IFD - 默认开启 chat template，可通过 --no_chat_template 切换回裸拼接 - chat template 缺失时给出 RuntimeError 提示	2026-06-18 16:35:05 +08:00
ViperEkura	a4e5a8c81c	feat: 新增 WSD 学习率调度器 - 支持 Warmup-Stable-Decay 三段式调度 - stable 阶段保持最高 lr，decay 阶段 sqrt 衰减 - 适用于持续预训练、SFT、RLHF 场景	2026-06-18 15:55:15 +08:00
ViperEkura	3e234c46f6	fix: 使用 threading.Event 替代裸 bool，补全公共 API - scheduler 停止信号改用 threading.Event，跨解释器安全 - 移除 _fatal_error 和 check_health，异常仅用 logger.error 记录 - 补全 astrai/__init__.py，暴露所有主要模块	2026-06-18 15:38:35 +08:00
ViperEkura	7a04b1f8ce	docs: replace shields.io endpoint badges with github/ direct badges - Switch stars/forks/release to github/ endpoints to avoid pool exhaustion - Add CI workflow badge for tests.yml - Delete update-badges.yml (no longer needed) - Remove remote gh-pages branch	2026-06-18 15:09:51 +08:00