use crate::write::{PackedWeight, WriteOptions, write_package};
use anyhow::Result;
use rlx_ir::Graph;
use std::path::Path;
#[derive(Debug, Clone)]
pub struct BakeWeight {
pub name: String,
pub shape: Vec<usize>,
pub encoding: String,
pub data: Vec<u8>,
}
impl BakeWeight {
fn layout_for_encoding(encoding: &str) -> &'static str {
if encoding.starts_with("gguf_") {
"bt_nk"
} else {
"row_major"
}
}
}
pub fn package_from_bake(
out: impl AsRef<Path>,
graph: &Graph,
weights: &[BakeWeight],
mut opts: WriteOptions,
) -> Result<()> {
if opts.name == "model" {
if !graph.name.is_empty() {
opts.name = graph.name.clone();
}
}
if opts.producer.is_none() {
opts.producer = Some("rlx-bake".into());
}
let packed: Vec<PackedWeight> = weights
.iter()
.map(|w| PackedWeight {
name: w.name.clone(),
shape: w.shape.clone(),
scheme: w.encoding.clone(),
layout: BakeWeight::layout_for_encoding(&w.encoding).into(),
data: w.data.clone(),
rank: None,
tier: crate::tier::StorageTier::Hot,
})
.collect();
write_package(out, graph, &packed, &opts)
}