# MiniCPM5-1B-rust

基于 [Burn](https://burn.dev/) 深度学习框架实现的 MiniCPM5-1B 大语言模型推理库，纯 Rust 编写，支持 GPU 加速。

## 模型架构参数

| 参数 | 值 |
|------|-----|
| 层数 | 24 |
| 注意力机制 | GQA (16 Query / 2 KV) |
| head_dim | 128 |
| hidden_size | 1536 |
| intermediate_size | 8960 |
| vocab_size | 151936 |
| RoPE theta | 5000000 |
| 归一化 | RMSNorm |
| 激活函数 | SiLU (SwiGLU) |
| KV Cache | 支持 |

## Workspace 结构

项目采用 Cargo workspace 多成员结构，包含三个 crate：

### minicpm-core

核心模型定义，包含：

- `LlamaConfig` — 模型配置（从 `config.json` 加载）
- `LlamaForCausalLM` — 因果语言模型主体
- `LlamaKVCache` — KV 缓存结构
- `EosTokenId` — EOS token 标识（支持单个或多个）
- 各模块实现：Attention、Decoder、FFN、RMSNorm、RoPE

### minicpm-convert

模型格式转换工具，负责将 HuggingFace safetensors 格式的权重转换为 Burn MPK 格式：

- `export_model()` — 完整转换流程（加载 safetensors → 构建模型 → 导出 MPK）
- 支持 BF16 / F16 / F32 精度输入
- 支持 `tie_word_embeddings` 权重共享

### minicpm-inference

推理功能封装，提供高层 API：

- `MiniCPM` — 高层封装，整合模型 + tokenizer + 推理逻辑
- `TokenizerWrapper` — tokenizer 封装，支持 MiniCPM5 chat template
- `GenerationConfig` — 生成配置（max_new_tokens、temperature、top_p）
- `generate_with_cache()` — 带 KV Cache 的自回归生成
- 支持 temperature 采样、top-p 采样、greedy 解码

## 快速开始

### 1. 转换模型

首先将 HuggingFace 格式的 MiniCPM5-1B 模型转换为 Burn MPK 格式：

```rust
use minicpm_convert::export_model;
use burn::backend::Wgpu;
use std::path::Path;

fn main() -> anyhow::Result<()> {
    let device = Default::default();

    export_model::<Wgpu>(
        Path::new("MiniCPM5-1B/model.safetensors"),
        Path::new("MiniCPM5-1B/config.json"),
        Path::new("MiniCPM5-1B-burn"),
        &device,
    )?;

    Ok(())
}
```

转换完成后，`MiniCPM5-1B-burn/` 目录下会生成 `model.mp` 文件。

### 2. 推理使用

使用 `MiniCPM` 高层 API 进行文本生成：

```rust
use minicpm_inference::{MiniCPM, GenerationConfig};
use burn::backend::Wgpu;

fn main() -> anyhow::Result<()> {
    let device = Default::default();

    // 加载模型
    let model = MiniCPM::<Wgpu>::load(
        "MiniCPM5-1B-burn/model",
        "MiniCPM5-1B/config.json",
        "MiniCPM5-1B/tokenizer.json",
        &device,
    )?;

    // 生成配置
    let gen_config = GenerationConfig {
        max_new_tokens: Some(512),
        temperature: 0.7,
        top_p: 0.8,
    };

    // 生成回答（think = true 启用思考模式）
    let response = model.generate("用 Rust 写一个 Hello World", true, &gen_config)?;

    println!("{}", response);

    Ok(())
}
```

## 移植到其他项目

### 依赖配置

在你的 `Cargo.toml` 中添加：

```toml
[dependencies]
minicpm-inference = { path = "../MiniCPM5-1B-rust/crates/minicpm-inference" }
burn = { version = "0.21", features = ["std", "wgpu"] }
anyhow = "1.0"
```

> 也可以根据需要选择其他 backend（如 `tch-gpu`、`cuda` 等）。

### 最小示例

```rust
use minicpm_inference::{MiniCPM, GenerationConfig};
use burn::backend::Wgpu;

fn main() -> anyhow::Result<()> {
    let device = Default::default();

    let model = MiniCPM::<Wgpu>::load(
        "model",
        "config.json",
        "tokenizer.json",
        &device,
    )?;

    let config = GenerationConfig {
        max_new_tokens: Some(256),
        temperature: 1.0,
        top_p: 1.0,
    };

    let output = model.generate("你好", false, &config)?;
    println!("{}", output);

    Ok(())
}
```

## 模型文件准备

`MiniCPM5-1B/` 目录需要包含以下文件：

```
MiniCPM5-1B/
├── config.json        # 模型配置文件
├── model.safetensors  # safetensors 格式权重
└── tokenizer.json     # tokenizer 配置
```

转换后 Burn 模型目录结构：

```
MiniCPM5-1B-burn/
└── model.mp           # Burn MPK 格式权重
```

## 性能说明

| Backend | 设备 | 速度 | 显存占用 |
|---------|------|------|----------|
| WGPU (Vulkan) | RTX 4060 | ~10 tokens/s | ~4 GB (F32) |

> 以上数据仅供参考，实际性能因硬件配置、生成长度等因素而异。

## 许可证

MIT