brain2qwerty 0.0.1

//! torchaudio Conformer (inference, dropout disabled).

use crate::tensor::Tensor;

pub struct FeedForward {
    pub ln_w: Tensor,
    pub ln_b: Tensor,
    pub w1: Tensor,
    pub b1: Tensor,
    pub w2: Tensor,
    pub b2: Tensor,
}

impl FeedForward {
    pub fn forward(&self, x: &Tensor) -> Tensor {
        x.layer_norm(&self.ln_w, &self.ln_b, 1e-5)
            .linear(&self.w1, Some(&self.b1))
            .silu()
            .linear(&self.w2, Some(&self.b2))
    }
}

pub struct ConvModule {
    pub ln_w: Tensor,
    pub ln_b: Tensor,
    pub p1_w: Tensor,
    pub p1_b: Tensor,
    pub dw_w: Tensor,
    pub dw_b: Tensor,
    pub gn_w: Tensor,
    pub gn_b: Tensor,
    pub p2_w: Tensor,
    pub p2_b: Tensor,
    pub kernel: usize,
    pub use_group_norm: bool,
}

impl ConvModule {
    pub fn forward(&self, input: &Tensor) -> Tensor {
        let x = input.layer_norm(&self.ln_w, &self.ln_b, 1e-5);
        let x_bct = x.transpose(&[0, 2, 1]);
        let y = x_bct.conv1d(&self.p1_w, Some(&self.p1_b), 1, 0, 1, 1);
        let half = y.shape[1] / 2;
        let y0 = y.narrow(1, 0, half);
        let y1 = y.narrow(1, half, half);
        let glu = y0.mul(&y1.sigmoid());
        let pad = self.kernel / 2;
        let mut dw = glu.conv1d(&self.dw_w, Some(&self.dw_b), 1, pad, 1, glu.shape[1]);
        if self.use_group_norm {
            dw = dw.group_norm(&self.gn_w, &self.gn_b, 1, 1e-5);
        } else {
            dw = dw.batch_norm1d(&self.gn_w, &self.gn_b, 1e-5);
        }
        dw.silu()
            .conv1d(&self.p2_w, Some(&self.p2_b), 1, 0, 1, 1)
            .transpose(&[0, 2, 1])
    }
}

pub struct ConformerLayer {
    pub ffn1: FeedForward,
    pub sa_ln_w: Tensor,
    pub sa_ln_b: Tensor,
    pub in_proj_w: Tensor,
    pub in_proj_b: Tensor,
    pub out_proj_w: Tensor,
    pub out_proj_b: Tensor,
    pub conv: ConvModule,
    pub ffn2: FeedForward,
    pub final_ln_w: Tensor,
    pub final_ln_b: Tensor,
    pub num_heads: usize,
    pub convolution_first: bool,
}

impl ConformerLayer {
    pub fn forward(&self, input: &Tensor) -> Tensor {
        let mut x = input.clone();
        let residual = x.clone();
        x = self.ffn1.forward(&x);
        x = x.mul_scalar(0.5).add(&residual);

        if self.convolution_first {
            x = self.apply_conv(&x);
        }

        let residual = x.clone();
        x = x.layer_norm(&self.sa_ln_w, &self.sa_ln_b, 1e-5);
        x = self.self_attn(&x);
        x = x.add(&residual);

        if !self.convolution_first {
            x = self.apply_conv(&x);
        }

        let residual = x.clone();
        x = self.ffn2.forward(&x);
        x = x.mul_scalar(0.5).add(&residual);
        x.layer_norm(&self.final_ln_w, &self.final_ln_b, 1e-5)
    }

    fn apply_conv(&self, input: &Tensor) -> Tensor {
        let residual = input.clone();
        self.conv.forward(input).add(&residual)
    }

    pub fn forward_self_attn(&self, x_normalized: &Tensor) -> Tensor {
        self.self_attn(x_normalized)
    }

    fn self_attn(&self, x: &Tensor) -> Tensor {
        let (b, t, d) = (x.shape[0], x.shape[1], x.shape[2]);
        let nh = self.num_heads;
        let dh = d / nh;
        let qkv = x.linear(&self.in_proj_w, Some(&self.in_proj_b));
        let mut q = Tensor::zeros(&[b, t, d]);
        let mut k = Tensor::zeros(&[b, t, d]);
        let mut v = Tensor::zeros(&[b, t, d]);
        for bi in 0..b {
            for ti in 0..t {
                let base = (bi * t + ti) * (3 * d);
                for j in 0..d {
                    q.data[(bi * t + ti) * d + j] = qkv.data[base + j];
                    k.data[(bi * t + ti) * d + j] = qkv.data[base + d + j];
                    v.data[(bi * t + ti) * d + j] = qkv.data[base + 2 * d + j];
                }
            }
        }
        let scale = (dh as f32).sqrt().recip();
        let mut out = Tensor::zeros(&[b, t, d]);
        for bi in 0..b {
            for hi in 0..nh {
                for ti in 0..t {
                    let mut scores = vec![0.0f32; t];
                    for tj in 0..t {
                        let mut dot = 0.0f32;
                        for j in 0..dh {
                            let q_idx = (bi * t + ti) * d + hi * dh + j;
                            let k_idx = (bi * t + tj) * d + hi * dh + j;
                            dot += q.data[q_idx] * k.data[k_idx];
                        }
                        scores[tj] = dot * scale;
                    }
                    let max_s = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                    let mut sum = 0.0f32;
                    for s in &mut scores {
                        *s = (*s - max_s).exp();
                        sum += *s;
                    }
                    for s in &mut scores {
                        *s /= sum;
                    }
                    for j in 0..dh {
                        let mut acc = 0.0f32;
                        for tj in 0..t {
                            let v_idx = (bi * t + tj) * d + hi * dh + j;
                            acc += scores[tj] * v.data[v_idx];
                        }
                        out.data[(bi * t + ti) * d + hi * dh + j] = acc;
                    }
                }
            }
        }
        out.linear(&self.out_proj_w, Some(&self.out_proj_b))
    }
}

pub struct ConformerStack {
    pub layers: Vec<ConformerLayer>,
}

impl ConformerStack {
    pub fn forward(&self, input: &Tensor) -> Tensor {
        let mut x = input.clone();
        for layer in &self.layers {
            x = layer.forward(&x);
        }
        x
    }
}