Lumotia/crates/core/src/model_registry.rs

use std::sync::LazyLock;

use crate::types::{Megabytes, ModelId};

/// Which inference backend a model uses.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Engine {
    Whisper,
    Parakeet,
    Moonshine,
}

/// Qualitative speed classification.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpeedTier {
    Instant,
    Fast,
    Moderate,
    Slow,
}

/// Qualitative accuracy classification.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccuracyTier {
    Excellent,
    Great,
    Good,
}

/// Language support scope.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LanguageSupport {
    EnglishOnly,
    Multilingual(u16),
}

/// File required for a model download.
#[derive(Debug, Clone)]
pub struct ModelFile {
    pub filename: &'static str,
    pub url: &'static str,
    pub size: Megabytes,
    /// SHA256 hex digest for integrity verification. None to skip check.
    pub sha256: Option<&'static str>,
}

/// All metadata for a single downloadable model.
/// This is pure data — no scoring logic lives here.
#[derive(Debug, Clone)]
pub struct ModelEntry {
    pub id: ModelId,
    pub engine: Engine,
    pub display_name: &'static str,
    pub disk_size: Megabytes,
    pub ram_required: Megabytes,
    pub speed_tier: SpeedTier,
    pub accuracy_tier: AccuracyTier,
    pub languages: LanguageSupport,
    pub files: Vec<ModelFile>,
    pub description: &'static str,
}

static ALL_MODELS: LazyLock<Vec<ModelEntry>> = LazyLock::new(|| {
    vec![
        ModelEntry {
            id: ModelId::new("parakeet-ctc-0.6b-int8"),
            engine: Engine::Parakeet,
            display_name: "Parakeet TDT 0.6B v2 (int8)",
            disk_size: Megabytes(650),
            ram_required: Megabytes(700),
            speed_tier: SpeedTier::Instant,
            accuracy_tier: AccuracyTier::Great,
            languages: LanguageSupport::EnglishOnly,
            files: vec![
                ModelFile {
                    filename: "encoder-model.int8.onnx",
                    url: "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v2-onnx/resolve/main/encoder-model.int8.onnx",
                    size: Megabytes(620),
                    sha256: None,
                },
                ModelFile {
                    filename: "decoder_joint-model.int8.onnx",
                    url: "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v2-onnx/resolve/main/decoder_joint-model.int8.onnx",
                    size: Megabytes(3),
                    sha256: None,
                },
                ModelFile {
                    filename: "nemo128.onnx",
                    url: "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v2-onnx/resolve/main/nemo128.onnx",
                    size: Megabytes(1),
                    sha256: None,
                },
                ModelFile {
                    filename: "vocab.txt",
                    url: "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v2-onnx/resolve/main/vocab.txt",
                    size: Megabytes(1),
                    sha256: None,
                },
            ],
            description: "Fastest local model — near-instant transcription",
        },
        ModelEntry {
            id: ModelId::new("whisper-tiny-en"),
            engine: Engine::Whisper,
            display_name: "Whisper Tiny (English)",
            disk_size: Megabytes(75),
            ram_required: Megabytes(390),
            speed_tier: SpeedTier::Fast,
            accuracy_tier: AccuracyTier::Good,
            languages: LanguageSupport::EnglishOnly,
            files: vec![ModelFile {
                filename: "ggml-tiny.en.bin",
                url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en.bin",
                size: Megabytes(75),
                sha256: None,
            }],
            description: "Bundled with app — works instantly",
        },
        ModelEntry {
            id: ModelId::new("whisper-base-en"),
            engine: Engine::Whisper,
            display_name: "Whisper Base (English)",
            disk_size: Megabytes(142),
            ram_required: Megabytes(500),
            speed_tier: SpeedTier::Fast,
            accuracy_tier: AccuracyTier::Good,
            languages: LanguageSupport::EnglishOnly,
            files: vec![ModelFile {
                filename: "ggml-base.en.bin",
                url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin",
                size: Megabytes(142),
                sha256: None,
            }],
            description: "Good balance of speed and accuracy",
        },
        ModelEntry {
            id: ModelId::new("whisper-small-en"),
            engine: Engine::Whisper,
            display_name: "Whisper Small (English)",
            disk_size: Megabytes(466),
            ram_required: Megabytes(1024),
            speed_tier: SpeedTier::Moderate,
            accuracy_tier: AccuracyTier::Great,
            languages: LanguageSupport::EnglishOnly,
            files: vec![ModelFile {
                filename: "ggml-small.en.bin",
                url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-small.en.bin",
                size: Megabytes(466),
                sha256: None,
            }],
            description: "Accuracy-first English transcription",
        },
        ModelEntry {
            id: ModelId::new("whisper-distil-small-en"),
            engine: Engine::Whisper,
            display_name: "Distil-Whisper Small (English)",
            disk_size: Megabytes(336),
            ram_required: Megabytes(900),
            speed_tier: SpeedTier::Fast,
            accuracy_tier: AccuracyTier::Great,
            languages: LanguageSupport::EnglishOnly,
            files: vec![ModelFile {
                filename: "ggml-distil-small.en.bin",
                url: "https://huggingface.co/distil-whisper/distil-small.en/resolve/main/ggml-distil-small.en.bin",
                size: Megabytes(336),
                sha256: None,
            }],
            description: "Small accuracy, ~6\u{00d7} faster — distilled variant",
        },
        ModelEntry {
            id: ModelId::new("whisper-medium-en"),
            engine: Engine::Whisper,
            display_name: "Whisper Medium (English)",
            disk_size: Megabytes(1500),
            ram_required: Megabytes(2600),
            speed_tier: SpeedTier::Slow,
            accuracy_tier: AccuracyTier::Excellent,
            languages: LanguageSupport::EnglishOnly,
            files: vec![ModelFile {
                filename: "ggml-medium.en.bin",
                url: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-medium.en.bin",
                size: Megabytes(1500),
                sha256: None,
            }],
            description: "Best Whisper accuracy — needs 4+ GB RAM",
        },
        ModelEntry {
            id: ModelId::new("whisper-distil-large-v3"),
            engine: Engine::Whisper,
            display_name: "Distil-Whisper Large v3 (English)",
            disk_size: Megabytes(1550),
            ram_required: Megabytes(2800),
            speed_tier: SpeedTier::Moderate,
            accuracy_tier: AccuracyTier::Excellent,
            languages: LanguageSupport::EnglishOnly,
            files: vec![ModelFile {
                filename: "ggml-distil-large-v3.bin",
                url: "https://huggingface.co/distil-whisper/distil-large-v3-ggml/resolve/main/ggml-distil-large-v3.bin",
                size: Megabytes(1550),
                sha256: None,
            }],
            description: "Near large-v3 accuracy at ~6\u{00d7} the speed",
        },
    ]
});

/// Returns all known models. Pure data, no scoring.
pub fn all_models() -> &'static [ModelEntry] {
    &ALL_MODELS
}

/// Find a model by its ID.
pub fn find_model(id: &ModelId) -> Option<&'static ModelEntry> {
    ALL_MODELS.iter().find(|m| &m.id == id)
}