brain2qwerty 0.0.1

//! LoRA weight merging for inference (PEFT-style).

use crate::weights::{ParamBuf, WeightStore};

pub fn merge_lora_into_base(
    store: &mut WeightStore,
    rank: usize,
    alpha: usize,
    target_modules: &[String],
) -> anyhow::Result<()> {
    let scale = alpha as f32 / rank as f32;
    let keys: Vec<String> = store.keys().cloned().collect();
    for key in keys {
        if !key.contains(".lora_A.") && !key.contains("lora_A") {
            continue;
        }
        let base_key = key
            .replace(".lora_A.default.weight", ".weight")
            .replace(".lora_A.weight", ".weight")
            .replace("lora_A", "");
        let b_key = key.replace("lora_A", "lora_B");
        if !store.contains_key(&b_key) {
            continue;
        }
        let module_ok = target_modules.iter().any(|m| key.contains(m));
        if !module_ok {
            continue;
        }
        let a = store.get(&key).cloned();
        let b = store.get(&b_key).cloned();
        if let (Some(a), Some(b)) = (a, b) {
            if let Some(base) = store.get_mut(&base_key) {
                apply_lora_delta(base, &a, &b, scale);
                store.remove(&key);
                store.remove(&b_key);
            }
        }
    }
    Ok(())
}

fn apply_lora_delta(base: &mut ParamBuf, a: &ParamBuf, b: &ParamBuf, scale: f32) {
    let rank = a.shape[0].min(a.shape.get(1).copied().unwrap_or(0));
    let (out_d, in_d) = (base.shape[0], base.shape[1]);
    for o in 0..out_d {
        for i in 0..in_d {
            let mut delta = 0.0f32;
            for r in 0..rank {
                let a_v = if a.shape.len() == 2 {
                    a.data[r * a.shape[1] + i.min(a.shape[1] - 1)]
                } else {
                    a.data[r]
                };
                let b_v = if b.shape.len() == 2 {
                    b.data[o * b.shape[1] + r.min(b.shape[1] - 1)]
                } else {
                    b.data[o]
                };
                delta += b_v * a_v;
            }
            base.data[o * in_d + i] += scale * delta;
        }
    }
}