brain2qwerty 0.0.1

//! Safetensors → RLX parameter map for the Conformer tail graph.

use std::collections::HashMap;

use crate::model_rlx::graph::EncoderSpec;
use crate::weights::{load_safetensors, ParamBuf, WeightStore};

pub type ParamMap = HashMap<String, ParamBuf>;

fn transpose(data: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut out = vec![0f32; data.len()];
    for r in 0..rows {
        for c in 0..cols {
            out[c * rows + r] = data[r * cols + c];
        }
    }
    out
}

fn take_linear_w(store: &WeightStore, key: &str) -> anyhow::Result<ParamBuf> {
    let p = store
        .get(key)
        .ok_or_else(|| anyhow::anyhow!("missing weight key: {key}"))?;
    anyhow::ensure!(
        p.shape.len() == 2,
        "Linear weight {key} must be 2-D, got {:?}",
        p.shape
    );
    let (out_d, in_d) = (p.shape[0], p.shape[1]);
    Ok(ParamBuf {
        data: transpose(&p.data, out_d, in_d),
        shape: vec![in_d, out_d],
    })
}

fn take_conv1d_as_conv2d(store: &WeightStore, key: &str) -> anyhow::Result<ParamBuf> {
    let p = store
        .get(key)
        .ok_or_else(|| anyhow::anyhow!("missing weight key: {key}"))?;
    anyhow::ensure!(p.shape.len() == 3, "conv1d weight {key} must be 3-D");
    let (co, ci, kw) = (p.shape[0], p.shape[1], p.shape[2]);
    // Map Conv1d O×I×K to Conv2d O×I×1×K (time on width).
    Ok(ParamBuf {
        data: p.data.clone(),
        shape: vec![co, ci, 1, kw],
    })
}

pub fn build_rlx_params(raw: &mut WeightStore, spec: &EncoderSpec) -> anyhow::Result<ParamMap> {
    let mut params = ParamMap::new();

    if spec.aux {
        copy_param(raw, &mut params, "shared_layer_norm.weight")?;
        copy_param(raw, &mut params, "shared_layer_norm.bias")?;
    }
    params.insert(
        "output_layer.weight".into(),
        take_linear_w(raw, "output_layer.weight")?,
    );
    copy_param(raw, &mut params, "output_layer.bias")?;

    for layer in 0..spec.num_layers {
        let p = format!("transformer.conformer_layers.{layer}");
        for suffix in [
            "ffn1.sequential.0.weight",
            "ffn1.sequential.0.bias",
            "ffn1.sequential.1.bias",
            "ffn1.sequential.4.bias",
            "self_attn_layer_norm.weight",
            "self_attn_layer_norm.bias",
            "ffn2.sequential.0.weight",
            "ffn2.sequential.0.bias",
            "ffn2.sequential.1.bias",
            "ffn2.sequential.4.bias",
            "final_layer_norm.weight",
            "final_layer_norm.bias",
        ] {
            copy_param(raw, &mut params, &format!("{p}.{suffix}"))?;
        }
        for suffix in [
            "ffn1.sequential.1.weight",
            "ffn1.sequential.4.weight",
            "ffn2.sequential.1.weight",
            "ffn2.sequential.4.weight",
        ] {
            params.insert(
                format!("{p}.{suffix}"),
                take_linear_w(raw, &format!("{p}.{suffix}"))?,
            );
        }
        params.insert(
            format!("{p}.self_attn.in_proj.weight"),
            take_linear_w(raw, &format!("{p}.self_attn.in_proj_weight"))?,
        );
        params.insert(
            format!("{p}.self_attn.out_proj.weight"),
            take_linear_w(raw, &format!("{p}.self_attn.out_proj.weight"))?,
        );
        copy_param(raw, &mut params, &format!("{p}.self_attn.in_proj_bias"))?;
        params.insert(
            format!("{p}.self_attn.in_proj.bias"),
            raw.get(&format!("{p}.self_attn.in_proj_bias"))
                .cloned()
                .ok_or_else(|| anyhow::anyhow!("missing in_proj bias"))?,
        );
        copy_param(raw, &mut params, &format!("{p}.self_attn.out_proj.bias"))?;
        let cp = format!("transformer.conformer_layers.{layer}.conv_module");
        for suffix in [
            "layer_norm.weight",
            "layer_norm.bias",
            "sequential.0.bias",
            "sequential.2.bias",
            "sequential.3.weight",
            "sequential.3.bias",
            "sequential.5.bias",
        ] {
            copy_param(raw, &mut params, &format!("{cp}.{suffix}"))?;
        }
        for suffix in [
            "sequential.0.weight",
            "sequential.2.weight",
            "sequential.5.weight",
        ] {
            params.insert(
                format!("{cp}.{suffix}"),
                take_conv1d_as_conv2d(raw, &format!("{cp}.{suffix}"))?,
            );
        }
    }
    Ok(params)
}

fn copy_param(raw: &WeightStore, params: &mut ParamMap, key: &str) -> anyhow::Result<()> {
    let p = raw
        .get(key)
        .ok_or_else(|| anyhow::anyhow!("missing {key}"))?;
    params.insert(key.to_string(), p.clone());
    Ok(())
}

pub fn apply_params(compiled: &mut rlx::CompiledGraph, params: &ParamMap) {
    for (name, buf) in params {
        compiled.set_param(name, &buf.data);
    }
}

pub fn load_encoder_weights(path: &str) -> anyhow::Result<WeightStore> {
    load_safetensors(path)
}