47 lines
1.2 KiB
Rust
47 lines
1.2 KiB
Rust
pub mod format;
|
|
pub mod loader;
|
|
pub mod utils;
|
|
|
|
pub use format::{ExportTask, Format};
|
|
pub use loader::TensorStore;
|
|
|
|
use burn::tensor::backend::Backend;
|
|
use std::path::Path;
|
|
|
|
use minicpm_core::config::LlamaConfig;
|
|
|
|
/// 执行导出任务
|
|
pub fn run_export<B: Backend>(
|
|
safetensors_path: &Path,
|
|
config_path: &Path,
|
|
tokenizer_path: &Path,
|
|
tasks: &[ExportTask],
|
|
device: &B::Device,
|
|
) -> anyhow::Result<()> {
|
|
println!("开始转换 MiniCPM 模型...");
|
|
println!("源文件: {:?}", safetensors_path);
|
|
|
|
println!("加载配置文件...");
|
|
let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
|
|
|
|
println!("解析 safetensors 文件...");
|
|
let store = TensorStore::from_file(safetensors_path)?;
|
|
|
|
for task in tasks {
|
|
println!("\n--- 导出 {} ---", task.format.name());
|
|
let output_dir = Path::new(&task.output_dir);
|
|
|
|
match task.format {
|
|
Format::Full => {
|
|
format::full::export::<B>(&store, &config, output_dir, device)?;
|
|
}
|
|
}
|
|
|
|
std::fs::copy(config_path, output_dir.join("config.json"))?;
|
|
std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
|
|
}
|
|
|
|
println!("\n所有导出任务完成！");
|
|
Ok(())
|
|
}
|