inferencelayer 0.2.3

//! ARK-ASR-3B — the audio half of the current Open ASR leaderboard leader (4.76% mean WER).
//!
//! The 2026 leaders all have the same shape: an audio encoder, an LLM decoder, and a small adapter
//! between them. ARK is the cleanest instance of it, and the one this engine was already holding
//! both halves of:
//!
//! ```text
//!   16 kHz audio ──logmel(128)──▶ Whisper-large-v3 ENCODER  (src/whisper.rs, prefix
//!                                 `audio_encoder.whisper`)          [1500, 1280]
//!            ──LayerNorm──▶ ──merge 4 frames──▶ [375, 5120]
//!            ──Linear 5120→4096 ─gelu─ Linear 4096→2048──▶ [375, 2048]
//!                                              │
//!   "<|user|><|begin_of_audio|>" + <|audio|>×375 + "<|end_of_audio|><|assistant|>"
//!                                              │ spliced at audio_token_id
//!                                              ▼
//!                        Qwen2.5-3B decoder (`Arch::Qwen2`, 36L, d2048, GQA 16/2)
//! ```
//!
//! `ArkasrForConditionalGeneration` subclasses `Qwen2ForCausalLM` in the reference, and its
//! `model.layers.*` are Qwen2 tensor-for-tensor, so the decoder needs no new architecture here —
//! only `weights.rs` learning the name.
//!
//! Two details from the reference that a shape-compatible port would silently get wrong:
//!
//! * **The Whisper encoder's final `layer_norm` is replaced by Identity**
//!   (`self.whisper.layer_norm = nn.Identity()`), and ARK applies its OWN norm before merging.
//!   The checkpoint therefore has no `audio_encoder.whisper.layer_norm`, and applying ours as well
//!   would corrupt every embedding. [`crate::whisper::WhisperEncoder::load_embedded`] loads that
//!   norm only if it exists, so the absence is honoured rather than papered over.
//! * **The encoder dims live in a nested `whisper_config`** while the top level describes the LLM.
//!   Reading `hidden_size` from the top level builds a 2048-wide encoder for 1280-wide weights.

use std::path::Path;

use anyhow::{Context, Result};

use crate::cpu_gemm::{PackedWeight, gemm_packed};
use crate::weights::LazySt;
use crate::whisper::WhisperEncoder;
use crate::whisper_gpu::WhisperEncoderGpu;

/// Geometry, read from the checkpoint's own config.
#[derive(Debug, Clone)]
pub struct ArkConfig {
    /// Whisper encoder width (1280 for large-v3).
    pub audio_dim: usize,
    /// LLM width (2048).
    pub hidden_size: usize,
    /// Encoder frames folded into one LLM token.
    pub merge_factor: usize,
    /// Placeholder token the audio embeddings are spliced onto.
    pub audio_token_id: u32,
    /// Encoder frames per 30 s window (1500).
    pub max_whisper_length: usize,
}

impl ArkConfig {
    pub fn load(dir: &Path) -> Result<Self> {
        let v: serde_json::Value =
            serde_json::from_slice(&std::fs::read(dir.join("config.json"))?).context("config")?;
        let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        let audio_dim = v
            .get("whisper_config")
            .and_then(|w| w.get("d_model"))
            .and_then(|x| x.as_u64())
            .context("config.whisper_config.d_model")? as usize;
        Ok(Self {
            audio_dim,
            hidden_size: g("hidden_size"),
            merge_factor: g("merge_factor").max(1),
            audio_token_id: g("audio_token_id") as u32,
            max_whisper_length: if g("max_whisper_length") == 0 {
                1500
            } else {
                g("max_whisper_length")
            },
        })
    }
}

/// `y = x·Wᵀ + b`, row-major `[out, in]`.
struct Linear {
    packed: PackedWeight,
    bias: Vec<f32>,
    out: usize,
    inp: usize,
}

impl Linear {
    fn load(st: &LazySt, prefix: &str, out: usize, inp: usize) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        anyhow::ensure!(
            w.len() == out * inp,
            "{prefix}.weight has {} elements, expected {out}×{inp}",
            w.len()
        );
        Ok(Self {
            packed: PackedWeight::new(&w, out, inp),
            bias: st.tensor_f32(&format!("{prefix}.bias"))?,
            out,
            inp,
        })
    }

    fn forward(&self, x: &[f32], rows: usize) -> Vec<f32> {
        debug_assert_eq!(x.len(), rows * self.inp);
        let mut y = vec![0f32; rows * self.out];
        gemm_packed(&mut y, x, &self.packed, rows, Some(&self.bias));
        y
    }
}

/// Affine LayerNorm over the last dimension.
struct LayerNorm {
    w: Vec<f32>,
    b: Vec<f32>,
    eps: f32,
}

impl LayerNorm {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        Ok(Self {
            w: st.tensor_f32(&format!("{prefix}.weight"))?,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
            eps: 1e-5,
        })
    }

    fn forward(&self, x: &mut [f32]) {
        let d = self.w.len();
        for row in x.chunks_mut(d) {
            let mean = row.iter().sum::<f32>() / d as f32;
            let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / d as f32;
            let inv = 1.0 / (var + self.eps).sqrt();
            for (i, v) in row.iter_mut().enumerate() {
                *v = (*v - mean) * inv * self.w[i] + self.b[i];
            }
        }
    }
}

/// Exact-erf GELU — HF's `"gelu"`, which is what `mlp_adapter_act` names.
fn gelu(x: f32) -> f32 {
    0.5 * x * (1.0 + libm::erff(x * std::f32::consts::FRAC_1_SQRT_2))
}

/// Audio → LLM embeddings. The decoder is loaded separately through the engine's Qwen2 path.
/// Where the 32-layer Whisper-large-v3 encoder runs. It is ~95% of the audio cost, so the GPU
/// path is the difference between ~30 s and a few seconds per clip.
enum EncoderPath {
    Cpu(Box<WhisperEncoder>),
    Gpu(Box<WhisperEncoderGpu>),
}

pub struct ArkAudioTower {
    pub cfg: ArkConfig,
    encoder: EncoderPath,
    norm: LayerNorm,
    proj_in: Linear,
    proj_out: Linear,
}

impl ArkAudioTower {
    /// CPU encoder.
    pub fn load(dir: &Path) -> Result<Self> {
        Self::build(dir, None)
    }

    /// GPU encoder — same numerics band as the rest of the engine's wgpu ports (cos ~1.0 vs CPU),
    /// and the only reason this model is usable interactively.
    pub fn load_on(ctx: &crate::GpuCtx, dir: &Path) -> Result<Self> {
        Self::build(dir, Some(ctx))
    }

    fn build(dir: &Path, ctx: Option<&crate::GpuCtx>) -> Result<Self> {
        let cfg = ArkConfig::load(dir)?;
        // The encoder's dims come from `whisper_config`, and its final layer_norm is Identity
        // here — both handled by `load_embedded`.
        let encoder =
            WhisperEncoder::load_embedded(dir, "audio_encoder.whisper", Some("whisper_config"))?;
        let encoder = match ctx {
            Some(ctx) => EncoderPath::Gpu(Box::new(WhisperEncoderGpu::new(
                ctx,
                encoder,
                cfg.max_whisper_length,
            )?)),
            None => EncoderPath::Cpu(Box::new(encoder)),
        };
        let st = LazySt::open(dir)?;
        let merged = cfg.audio_dim * cfg.merge_factor;
        Ok(Self {
            norm: LayerNorm::load(&st, "audio_encoder.layer_norm")?,
            // Sequential(Linear(audio*k → 2*hidden), gelu, Linear(2*hidden → hidden)).
            proj_in: Linear::load(&st, "audio_encoder.adapting.0", cfg.hidden_size * 2, merged)?,
            proj_out: Linear::load(
                &st,
                "audio_encoder.adapting.2",
                cfg.hidden_size,
                cfg.hidden_size * 2,
            )?,
            encoder,
            cfg,
        })
    }

    /// How many `<|audio|>` placeholders a clip of `mel_frames` needs — the reference's
    /// `calculate_audio_token_count`: the encoder halves the frame count, then `merge_factor`
    /// frames fold into one token.
    pub fn audio_token_count(&self, mel_frames: usize) -> usize {
        // `(x + 1) / 2` in the reference — the same thing, said the way clippy prefers.
        (mel_frames.div_ceil(2) / self.cfg.merge_factor).max(1)
    }

    /// 16 kHz mono → `[tokens, hidden_size]` row-major, ready to splice into the decoder's
    /// input embeddings.
    pub fn encode(&self, ctx: Option<&crate::GpuCtx>, samples_16k: &[f32]) -> Result<Vec<f32>> {
        // No final Whisper norm on either path — ARK replaced it with Identity.
        let enc = match &self.encoder {
            EncoderPath::Cpu(e) => e.encode(samples_16k),
            EncoderPath::Gpu(g) => {
                let ctx =
                    ctx.context("this tower was built with a GPU encoder; pass its GpuCtx")?;
                let mel = g.cpu().logmel(samples_16k);
                g.forward(ctx, &mel)?
            }
        };
        let (d, k) = (self.cfg.audio_dim, self.cfg.merge_factor);
        let mut h = enc;
        self.norm.forward(&mut h);

        // Fold `k` consecutive frames into one row; drop the ragged tail exactly as the reference
        // does (it truncates rather than padding whenever a whole group remains).
        let frames = h.len() / d;
        let groups = (frames / k).max(1);
        h.resize(groups * k * d, 0.0);

        let mid = self.proj_in.forward(&h, groups);
        let act: Vec<f32> = mid.into_iter().map(gelu).collect();
        Ok(self.proj_out.forward(&act, groups))
    }
}

// =================================================================================================
// the whole model: audio tower + Qwen2.5 decoder
// =================================================================================================

/// The instruction the reference's own example uses after the audio part.
#[cfg(feature = "cli")]
pub const DEFAULT_INSTRUCTION: &str = "Please transcribe this audio.";

/// ARK-ASR-3B end to end.
///
/// The decoder is the engine's ordinary Qwen2 path (`Arch::Qwen2`) — the reference subclasses
/// `Qwen2ForCausalLM`, so nothing about it is special. The only unusual step is the prefill: text
/// tokens go through `forward_only`, and the 375 `<|audio|>` placeholders are replaced by the
/// tower's embeddings through `forward_from_embeds_only`, which is the same splice the reference
/// does with `masked_scatter`.
#[cfg(feature = "cli")]
pub struct ArkAsr {
    pub tower: ArkAudioTower,
    gpu: crate::Lfm2Gpu,
    tok: tokenizers::Tokenizer,
    dir: std::path::PathBuf,
    eos: u32,
}

#[cfg(feature = "cli")]
impl ArkAsr {
    pub fn load(ctx: &crate::GpuCtx, dir: &Path) -> Result<Self> {
        let tower = ArkAudioTower::load_on(ctx, dir)?;
        let w = crate::Weights::load(ctx, dir).context("load ARK decoder (Arch::Qwen2)")?;
        anyhow::ensure!(
            w.cfg.hidden == tower.cfg.hidden_size,
            "decoder hidden {} != adapter output {}",
            w.cfg.hidden,
            tower.cfg.hidden_size
        );
        let eos =
            serde_json::from_slice::<serde_json::Value>(&std::fs::read(dir.join("config.json"))?)?
                .get("eos_token_id")
                .and_then(|v| v.as_u64())
                .unwrap_or(151645) as u32;
        Ok(Self {
            gpu: crate::Lfm2Gpu::new(ctx, w),
            tok: tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
                .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?,
            dir: dir.to_path_buf(),
            eos,
            tower,
        })
    }

    fn ids(&self, text: &str) -> Result<Vec<u32>> {
        Ok(self
            .tok
            .encode(text, false)
            .map_err(|e| anyhow::anyhow!("encode: {e}"))?
            .get_ids()
            .to_vec())
    }

    /// 16 kHz mono → transcript.
    ///
    /// Prompt shape is the reference's: `<|user|><|begin_of_audio|>` + one `<|audio|>` per audio
    /// embedding + `<|end_of_audio|>` + the instruction + `<|assistant|>`. The placeholder count
    /// must equal the embedding count or the splice slides out of alignment, so it is derived from
    /// the tower rather than assumed.
    pub fn transcribe(
        &self,
        ctx: &crate::GpuCtx,
        samples_16k: &[f32],
        max_new: usize,
    ) -> Result<String> {
        self.transcribe_with(ctx, samples_16k, DEFAULT_INSTRUCTION, max_new)
    }

    /// As [`Self::transcribe`], with the instruction that follows the audio.
    ///
    /// This model is an LLM, so the instruction is not decoration: with it omitted, English still
    /// transcribed but French returned a bare ".". The reference's own example passes
    /// "Please transcribe this audio." after the audio part, and that is what makes the
    /// non-English languages it advertises (zh/de/ja/fr/ko/es/pl/it) actually work.
    pub fn transcribe_with(
        &self,
        ctx: &crate::GpuCtx,
        samples_16k: &[f32],
        instruction: &str,
        max_new: usize,
    ) -> Result<String> {
        // Each utterance is independent: clear the decoder state so nothing leaks from the last
        // clip. Without this the second call inherits the first's session and can emit a bare
        // "." and stop.
        self.gpu.reset(ctx);
        let hidden = self.tower.cfg.hidden_size;
        let audio = self.tower.encode(Some(ctx), samples_16k)?;
        let n_audio = audio.len() / hidden;

        let pre = self.ids("<|user|><|begin_of_audio|>")?;
        let post = self.ids(&format!("<|end_of_audio|>{instruction}<|assistant|>"))?;

        // Prefill. Text goes through the token path; each audio slot through the embedding path.
        let mut pos = 0usize;
        for &t in &pre {
            self.gpu.forward_only(ctx, t, pos);
            pos += 1;
        }
        for i in 0..n_audio {
            self.gpu
                .forward_from_embeds_only(ctx, &audio[i * hidden..(i + 1) * hidden], pos)?;
            pos += 1;
        }
        for &t in &post[..post.len().saturating_sub(1)] {
            self.gpu.forward_only(ctx, t, pos);
            pos += 1;
        }

        // Greedy decode from the last prompt token.
        let mut next = *post.last().context("empty prompt tail")?;
        let mut out: Vec<u32> = Vec::new();
        for _ in 0..max_new {
            let logits = self.gpu.forward(ctx, next, pos)?;
            pos += 1;
            let (best, _) =
                logits
                    .iter()
                    .enumerate()
                    .fold((0usize, f32::NEG_INFINITY), |acc, (i, &v)| {
                        if v > acc.1 { (i, v) } else { acc }
                    });
            next = best as u32;
            if next == self.eos {
                break;
            }
            out.push(next);
        }
        let _ = &self.dir;
        Ok(self
            .tok
            .decode(&out, true)
            .map_err(|e| anyhow::anyhow!("decode: {e}"))?
            .trim()
            .to_string())
    }
}