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