chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user