Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,532 @@
+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());
+    }
+}