rustytorch/demos/rtx-model-zoo/src/registry.rs

use model_zoo_shared::{ModelCategory, ModelInfo};
use std::collections::HashMap;

pub struct ModelRegistry {
    models: HashMap<String, ModelInfo>,
}

impl ModelRegistry {
    pub fn new() -> Self {
        let mut models = HashMap::new();

        models.insert(
            "resnet50".to_string(),
            ModelInfo {
                id: "resnet50".to_string(),
                name: "ResNet-50".to_string(),
                description:
                    "50-layer residual network for image classification, trained on ImageNet-1K"
                        .to_string(),
                category: ModelCategory::ImageClassification,
                size_mb: 97.8,
                parameters: 25_600_000,
                accuracy_metric: "Top-1: 76.1%, Top-5: 92.9%".to_string(),
                download_url: "https://download.pytorch.org/models/resnet50-0676ba61.pth"
                    .to_string(),
                license: "BSD-3-Clause".to_string(),
            },
        );

        models.insert(
            "efficientnet_b0".to_string(),
            ModelInfo {
                id: "efficientnet_b0".to_string(),
                name: "EfficientNet-B0".to_string(),
                description: "Efficient convolutional neural network with compound scaling, optimized for mobile and edge devices".to_string(),
                category: ModelCategory::ImageClassification,
                size_mb: 20.5,
                parameters: 5_300_000,
                accuracy_metric: "Top-1: 77.1%, Top-5: 93.3%".to_string(),
                download_url: "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth".to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        models.insert(
            "yolov8n".to_string(),
            ModelInfo {
                id: "yolov8n".to_string(),
                name: "YOLOv8-Nano".to_string(),
                description: "Ultra-fast object detection model from Ultralytics, optimized for real-time inference on edge devices".to_string(),
                category: ModelCategory::ObjectDetection,
                size_mb: 6.2,
                parameters: 3_200_000,
                accuracy_metric: "mAP@50: 37.3%, mAP@50-95: 28.0%".to_string(),
                download_url: "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt".to_string(),
                license: "AGPL-3.0".to_string(),
            },
        );

        models.insert(
            "detr_resnet50".to_string(),
            ModelInfo {
                id: "detr_resnet50".to_string(),
                name: "DETR-ResNet50".to_string(),
                description: "DEtection TRansformer with ResNet-50 backbone, end-to-end object detection using transformers".to_string(),
                category: ModelCategory::ObjectDetection,
                size_mb: 159.0,
                parameters: 41_300_000,
                accuracy_metric: "mAP@50: 62.4%, mAP@50-95: 42.0%".to_string(),
                download_url: "https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth".to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        models.insert(
            "deeplabv3_resnet50".to_string(),
            ModelInfo {
                id: "deeplabv3_resnet50".to_string(),
                name: "DeepLabV3-ResNet50".to_string(),
                description:
                    "State-of-the-art semantic segmentation with atrous convolution and ASPP module"
                        .to_string(),
                category: ModelCategory::Segmentation,
                size_mb: 158.0,
                parameters: 39_600_000,
                accuracy_metric: "mIoU: 77.2% (PASCAL VOC 2012)".to_string(),
                download_url:
                    "https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth"
                        .to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        models.insert(
            "segformer_b0".to_string(),
            ModelInfo {
                id: "segformer_b0".to_string(),
                name: "SegFormer-B0".to_string(),
                description: "Efficient transformer-based semantic segmentation with hierarchical design".to_string(),
                category: ModelCategory::Segmentation,
                size_mb: 14.3,
                parameters: 3_700_000,
                accuracy_metric: "mIoU: 76.2% (ADE20K)".to_string(),
                download_url: "https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/pytorch_model.bin".to_string(),
                license: "NVIDIA".to_string(),
            },
        );

        models.insert(
            "gpt2_small".to_string(),
            ModelInfo {
                id: "gpt2_small".to_string(),
                name: "GPT-2 Small".to_string(),
                description: "117M parameter autoregressive language model for text generation"
                    .to_string(),
                category: ModelCategory::TextGeneration,
                size_mb: 548.0,
                parameters: 117_000_000,
                accuracy_metric: "Perplexity: 29.4 (WebText)".to_string(),
                download_url: "https://huggingface.co/gpt2/resolve/main/pytorch_model.bin"
                    .to_string(),
                license: "MIT".to_string(),
            },
        );

        models.insert(
            "bert_base_uncased".to_string(),
            ModelInfo {
                id: "bert_base_uncased".to_string(),
                name: "BERT-Base Uncased".to_string(),
                description: "Bidirectional transformer for natural language understanding tasks (12-layer, 768-hidden, 12-heads)".to_string(),
                category: ModelCategory::NLP,
                size_mb: 438.0,
                parameters: 110_000_000,
                accuracy_metric: "GLUE Score: 79.6".to_string(),
                download_url: "https://huggingface.co/bert-base-uncased/resolve/main/pytorch_model.bin".to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        models.insert(
            "distilbert_base".to_string(),
            ModelInfo {
                id: "distilbert_base".to_string(),
                name: "DistilBERT-Base".to_string(),
                description: "Distilled version of BERT with 40% fewer parameters, 60% faster while retaining 97% of BERT's performance".to_string(),
                category: ModelCategory::NLP,
                size_mb: 268.0,
                parameters: 66_000_000,
                accuracy_metric: "GLUE Score: 77.0".to_string(),
                download_url: "https://huggingface.co/distilbert-base-uncased/resolve/main/pytorch_model.bin".to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        models.insert(
            "wav2vec2_base".to_string(),
            ModelInfo {
                id: "wav2vec2_base".to_string(),
                name: "Wav2Vec2-Base".to_string(),
                description:
                    "Self-supervised speech recognition model trained on unlabeled audio data"
                        .to_string(),
                category: ModelCategory::SpeechRecognition,
                size_mb: 378.0,
                parameters: 95_000_000,
                accuracy_metric: "WER: 6.1% (LibriSpeech clean)".to_string(),
                download_url:
                    "https://huggingface.co/facebook/wav2vec2-base/resolve/main/pytorch_model.bin"
                        .to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        models.insert(
            "whisper_tiny".to_string(),
            ModelInfo {
                id: "whisper_tiny".to_string(),
                name: "Whisper Tiny".to_string(),
                description:
                    "Multilingual speech recognition model with 39M parameters, optimized for speed"
                        .to_string(),
                category: ModelCategory::SpeechRecognition,
                size_mb: 73.0,
                parameters: 39_000_000,
                accuracy_metric: "WER: 7.5% (LibriSpeech)".to_string(),
                download_url:
                    "https://huggingface.co/openai/whisper-tiny/resolve/main/pytorch_model.bin"
                        .to_string(),
                license: "MIT".to_string(),
            },
        );

        models.insert(
            "vit_base_patch16".to_string(),
            ModelInfo {
                id: "vit_base_patch16".to_string(),
                name: "ViT-Base-Patch16".to_string(),
                description: "Vision Transformer with 16x16 patches, pure transformer architecture for image classification".to_string(),
                category: ModelCategory::ImageClassification,
                size_mb: 346.0,
                parameters: 86_600_000,
                accuracy_metric: "Top-1: 84.5% (ImageNet-21K pretrained)".to_string(),
                download_url: "https://huggingface.co/google/vit-base-patch16-224/resolve/main/pytorch_model.bin".to_string(),
                license: "Apache-2.0".to_string(),
            },
        );

        Self { models }
    }

    pub fn get_model(&self, model_id: &str) -> Option<&ModelInfo> {
        self.models.get(model_id)
    }

    pub fn list_all_models(&self) -> Vec<&ModelInfo> {
        self.models.values().collect()
    }

    pub fn filter_by_category(&self, category: ModelCategory) -> Vec<&ModelInfo> {
        self.models
            .values()
            .filter(|model| model.category == category)
            .collect()
    }

    pub fn count_models(&self) -> usize {
        self.models.len()
    }

    pub fn has_model(&self, model_id: &str) -> bool {
        self.models.contains_key(model_id)
    }
}

impl Default for ModelRegistry {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_registry_creation() {
        let registry = ModelRegistry::new();
        assert!(registry.count_models() >= 12);
    }

    #[test]
    fn test_registry_get_resnet50() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("resnet50");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "resnet50");
        assert_eq!(model.name, "ResNet-50");
        assert_eq!(model.category, ModelCategory::ImageClassification);
        assert_eq!(model.size_mb, 97.8);
        assert_eq!(model.parameters, 25_600_000);
    }

    #[test]
    fn test_registry_get_efficientnet() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("efficientnet_b0");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "efficientnet_b0");
        assert_eq!(model.name, "EfficientNet-B0");
        assert_eq!(model.category, ModelCategory::ImageClassification);
    }

    #[test]
    fn test_registry_get_yolov8() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("yolov8n");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "yolov8n");
        assert_eq!(model.name, "YOLOv8-Nano");
        assert_eq!(model.category, ModelCategory::ObjectDetection);
    }

    #[test]
    fn test_registry_get_detr() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("detr_resnet50");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "detr_resnet50");
        assert_eq!(model.category, ModelCategory::ObjectDetection);
    }

    #[test]
    fn test_registry_get_deeplabv3() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("deeplabv3_resnet50");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "deeplabv3_resnet50");
        assert_eq!(model.category, ModelCategory::Segmentation);
    }

    #[test]
    fn test_registry_get_segformer() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("segformer_b0");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "segformer_b0");
        assert_eq!(model.category, ModelCategory::Segmentation);
    }

    #[test]
    fn test_registry_get_gpt2() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("gpt2_small");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "gpt2_small");
        assert_eq!(model.category, ModelCategory::TextGeneration);
    }

    #[test]
    fn test_registry_get_bert() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("bert_base_uncased");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "bert_base_uncased");
        assert_eq!(model.category, ModelCategory::NLP);
    }

    #[test]
    fn test_registry_get_distilbert() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("distilbert_base");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "distilbert_base");
        assert_eq!(model.category, ModelCategory::NLP);
    }

    #[test]
    fn test_registry_get_wav2vec2() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("wav2vec2_base");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "wav2vec2_base");
        assert_eq!(model.category, ModelCategory::SpeechRecognition);
    }

    #[test]
    fn test_registry_get_whisper() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("whisper_tiny");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "whisper_tiny");
        assert_eq!(model.category, ModelCategory::SpeechRecognition);
    }

    #[test]
    fn test_registry_get_vit() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("vit_base_patch16");
        assert!(model.is_some());
        let model = model.unwrap();
        assert_eq!(model.id, "vit_base_patch16");
        assert_eq!(model.category, ModelCategory::ImageClassification);
    }

    #[test]
    fn test_registry_get_nonexistent() {
        let registry = ModelRegistry::new();
        let model = registry.get_model("nonexistent_model");
        assert!(model.is_none());
    }

    #[test]
    fn test_registry_list_all_models() {
        let registry = ModelRegistry::new();
        let models = registry.list_all_models();
        assert!(models.len() >= 12);
    }

    #[test]
    fn test_registry_filter_image_classification() {
        let registry = ModelRegistry::new();
        let models = registry.filter_by_category(ModelCategory::ImageClassification);
        assert!(models.len() >= 3);
        for model in models {
            assert_eq!(model.category, ModelCategory::ImageClassification);
        }
    }

    #[test]
    fn test_registry_filter_object_detection() {
        let registry = ModelRegistry::new();
        let models = registry.filter_by_category(ModelCategory::ObjectDetection);
        assert!(models.len() >= 2);
        for model in models {
            assert_eq!(model.category, ModelCategory::ObjectDetection);
        }
    }

    #[test]
    fn test_registry_filter_segmentation() {
        let registry = ModelRegistry::new();
        let models = registry.filter_by_category(ModelCategory::Segmentation);
        assert!(models.len() >= 2);
        for model in models {
            assert_eq!(model.category, ModelCategory::Segmentation);
        }
    }

    #[test]
    fn test_registry_filter_text_generation() {
        let registry = ModelRegistry::new();
        let models = registry.filter_by_category(ModelCategory::TextGeneration);
        assert!(models.len() >= 1);
        for model in models {
            assert_eq!(model.category, ModelCategory::TextGeneration);
        }
    }

    #[test]
    fn test_registry_filter_nlp() {
        let registry = ModelRegistry::new();
        let models = registry.filter_by_category(ModelCategory::NLP);
        assert!(models.len() >= 2);
        for model in models {
            assert_eq!(model.category, ModelCategory::NLP);
        }
    }

    #[test]
    fn test_registry_filter_speech_recognition() {
        let registry = ModelRegistry::new();
        let models = registry.filter_by_category(ModelCategory::SpeechRecognition);
        assert!(models.len() >= 2);
        for model in models {
            assert_eq!(model.category, ModelCategory::SpeechRecognition);
        }
    }

    #[test]
    fn test_registry_count_models() {
        let registry = ModelRegistry::new();
        let count = registry.count_models();
        assert_eq!(count, 12);
    }

    #[test]
    fn test_registry_has_model() {
        let registry = ModelRegistry::new();
        assert!(registry.has_model("resnet50"));
        assert!(registry.has_model("yolov8n"));
        assert!(!registry.has_model("nonexistent_model"));
    }

    #[test]
    fn test_registry_default() {
        let registry = ModelRegistry::default();
        assert_eq!(registry.count_models(), 12);
    }

    #[test]
    fn test_all_models_have_valid_data() {
        let registry = ModelRegistry::new();
        let models = registry.list_all_models();
        for model in models {
            assert!(!model.id.is_empty());
            assert!(!model.name.is_empty());
            assert!(!model.description.is_empty());
            assert!(model.size_mb > 0.0);
            assert!(model.parameters > 0);
            assert!(!model.accuracy_metric.is_empty());
            assert!(!model.download_url.is_empty());
            assert!(!model.license.is_empty());
        }
    }

    #[test]
    fn test_model_ids_are_unique() {
        let registry = ModelRegistry::new();
        let models = registry.list_all_models();
        let mut ids = std::collections::HashSet::new();
        for model in models {
            assert!(
                ids.insert(model.id.clone()),
                "Duplicate model ID: {}",
                model.id
            );
        }
    }

    #[test]
    fn test_all_categories_represented() {
        let registry = ModelRegistry::new();
        assert!(
            !registry
                .filter_by_category(ModelCategory::ImageClassification)
                .is_empty()
        );
        assert!(
            !registry
                .filter_by_category(ModelCategory::ObjectDetection)
                .is_empty()
        );
        assert!(
            !registry
                .filter_by_category(ModelCategory::Segmentation)
                .is_empty()
        );
        assert!(
            !registry
                .filter_by_category(ModelCategory::TextGeneration)
                .is_empty()
        );
        assert!(
            !registry
                .filter_by_category(ModelCategory::SpeechRecognition)
                .is_empty()
        );
        assert!(!registry.filter_by_category(ModelCategory::NLP).is_empty());
    }
}