brain2qwerty 0.0.1

//! Brain2Qwerty V1 inference pipeline.

use crate::tensor::Tensor;
use crate::v1::config::V1Config;
use crate::v1::decode::greedy_decode;
use crate::v1::model::Brain2QwertyV1;

/// One sentence worth of keystroke-aligned MEG windows.
pub struct V1Input {
    /// Keystroke windows `(N, channels, time)`.
    pub neuros: Tensor,
    pub subject_ids: Vec<usize>,
    pub chan_pos: Option<Tensor>,
}

pub struct V1Output {
    pub pred_text: String,
    pub logits: Tensor,
}

/// End-to-end V1 decoder.
pub struct V1Pipeline {
    pub model: Brain2QwertyV1,
}

impl V1Pipeline {
    /// Tiny CI weights from `generate_v1_parity_refs.py`.
    pub fn from_tiny_weights(weights_path: &str) -> anyhow::Result<Self> {
        Ok(Self {
            model: Brain2QwertyV1::from_tiny_weights(weights_path)?,
        })
    }

    /// Load from `v1_config.yaml` and a safetensors bundle (encoder + transformer + linear).
    pub fn from_weights(config_path: &str, weights_path: &str) -> anyhow::Result<Self> {
        let cfg = V1Config::from_yaml(config_path)?;
        Ok(Self {
            model: Brain2QwertyV1::from_config_and_weights(cfg, weights_path)?,
        })
    }

    /// Decode one sentence and return greedy character predictions.
    pub fn run(&self, input: &V1Input) -> anyhow::Result<V1Output> {
        let logits =
            self.model
                .forward_sentence(&input.neuros, &input.subject_ids, input.chan_pos.as_ref());
        Ok(V1Output {
            pred_text: greedy_decode(&logits),
            logits,
        })
    }
}