18 lines
579 B
Rust
18 lines
579 B
Rust
|
use burn::module::Module;
|
||
|
|
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||
|
|
use burn::tensor::backend::Backend;
|
||
|
|
|
||
|
|
use minicpm_core::config::LlamaConfig;
|
||
|
|
use minicpm_core::model::LlamaForCausalLM;
|
||
|
|
|
||
|
|
pub fn load_model<B: Backend>(
|
||
|
|
model_path: &str,
|
||
|
|
config: &LlamaConfig,
|
||
|
|
device: &B::Device,
|
||
|
|
) -> anyhow::Result<LlamaForCausalLM<B>> {
|
||
|
|
let model = LlamaForCausalLM::<B>::new(config.clone(), device);
|
||
|
|
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
|
||
|
|
let model = model.load_file(model_path, &recorder, device)?;
|
||
|
|
Ok(model)
|
||
|
|
}
|