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(())
|
|||
|
|
}
|