32 lines
1022 B
Rust
32 lines
1022 B
Rust
use burn::module::Module;
|
||
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||
use burn::tensor::backend::Backend;
|
||
use std::path::Path;
|
||
|
||
use crate::loader::{load_into_model, TensorStore};
|
||
use minicpm_core::config::LlamaConfig;
|
||
use minicpm_core::model::LlamaForCausalLM;
|
||
|
||
/// 导出全精度模型（f32）
|
||
pub fn export<B: Backend>(
|
||
store: &TensorStore,
|
||
config: &LlamaConfig,
|
||
output_dir: &Path,
|
||
device: &B::Device,
|
||
) -> anyhow::Result<()> {
|
||
println!("创建模型结构...");
|
||
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||
|
||
println!("加载 safetensors 权重...");
|
||
let model = load_into_model(model, store, device);
|
||
|
||
println!("保存为 MPK 格式（全精度）...");
|
||
std::fs::create_dir_all(output_dir)?;
|
||
let output_path = output_dir.join("model");
|
||
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||
model.save_file(&output_path, &recorder)?;
|
||
|
||
println!("全精度模型已导出到: {:?}", output_dir);
|
||
Ok(())
|
||
}
|