Merge pull request #8 from marqo-ai/typed-model-properties

Expose ModelProperties as objects
marqo-ai · Aug 22, 2024 · fe893ff · fe893ff
2 parents 0689006 + bf6edf4
commit fe893ff
Show file tree

Hide file tree

Showing 14 changed files with 1,667 additions and 1,610 deletions.
diff --git a/src/marqo_commons/model_registry/model_properties_data/clip_properties.py b/src/marqo_commons/model_registry/model_properties_data/clip_properties.py
@@ -4,7 +4,8 @@
 This file defines properties for CLIP models. It is intended to be used in conjunction with the model registry
 and should not be used in isolation.
 """
-from marqo_commons.model_registry.model_properties_object import ModelProperties, VectorNumericType, Modality, ModelType
+from marqo_commons.model_registry.model_properties_object import ModelProperties, VectorNumericType, Modality, \
+    ModelType, T
 from typing import Dict, List
 from marqo_commons.model_registry.utils import convert_model_properties_to_dict
 
@@ -16,53 +17,56 @@ class ClipProperties(ModelProperties):
     type: ModelType = ModelType.clip
 
 
+    @classmethod
+    def get_all_model_properties_objects(cls) -> Dict[str, T]:
+        return {
+            'RN50': ClipProperties(
+                name="RN50",
+                dimensions=1024,
+                notes="CLIP resnet50",
+            ),
+            'RN101': ClipProperties(
+                name="RN101",
+                dimensions=512,
+                notes="CLIP resnet101",
+            ),
+            'RN50x4': ClipProperties(
+                name="RN50x4",
+                dimensions=640,
+                notes="CLIP resnet50x4",
+            ),
+            'RN50x16': ClipProperties(
+                name="RN50x16",
+                dimensions=768,
+                notes="CLIP resnet50x16",
+            ),
+            'RN50x64': ClipProperties(
+                name="RN50x64",
+                dimensions=1024,
+                notes="CLIP resnet50x64",
+            ),
+            'ViT-B/32': ClipProperties(
+                name="ViT-B/32",
+                dimensions=512,
+                notes="CLIP ViT-B/32",
+            ),
+            'ViT-B/16': ClipProperties(
+                name="ViT-B/16",
+                dimensions=512,
+                notes="CLIP ViT-B/16",
+            ),
+            'ViT-L/14': ClipProperties(
+                name="ViT-L/14",
+                dimensions=768,
+                notes="CLIP ViT-L/14",
+            ),
+            'ViT-L/14@336px': ClipProperties(
+                name="ViT-L/14@336px",
+                dimensions=768,
+                notes="CLIP ViT-L/14@336px",
+            ),
+        }
+
 @convert_model_properties_to_dict
 def _get_clip_properties() -> Dict:
-    CLIP_MODEL_PROPERTIES = {
-        'RN50': ClipProperties(
-            name="RN50",
-            dimensions=1024,
-            notes="CLIP resnet50",
-        ),
-        'RN101': ClipProperties(
-            name="RN101",
-            dimensions=512,
-            notes="CLIP resnet101",
-        ),
-        'RN50x4': ClipProperties(
-            name="RN50x4",
-            dimensions=640,
-            notes="CLIP resnet50x4",
-        ),
-        'RN50x16': ClipProperties(
-            name="RN50x16",
-            dimensions=768,
-            notes="CLIP resnet50x16",
-        ),
-        'RN50x64': ClipProperties(
-            name="RN50x64",
-            dimensions=1024,
-            notes="CLIP resnet50x64",
-        ),
-        'ViT-B/32': ClipProperties(
-            name="ViT-B/32",
-            dimensions=512,
-            notes="CLIP ViT-B/32",
-        ),
-        'ViT-B/16': ClipProperties(
-            name="ViT-B/16",
-            dimensions=512,
-            notes="CLIP ViT-B/16",
-        ),
-        'ViT-L/14': ClipProperties(
-            name="ViT-L/14",
-            dimensions=768,
-            notes="CLIP ViT-L/14",
-        ),
-        'ViT-L/14@336px': ClipProperties(
-            name="ViT-L/14@336px",
-            dimensions=768,
-            notes="CLIP ViT-L/14@336px",
-        ),
-    }
-    return CLIP_MODEL_PROPERTIES
+    return ClipProperties.get_all_model_properties_objects()
diff --git a/src/marqo_commons/model_registry/model_properties_data/fp16_clip_properties.py b/src/marqo_commons/model_registry/model_properties_data/fp16_clip_properties.py
@@ -6,7 +6,8 @@
 """
 from typing import Dict, List
 
-from marqo_commons.model_registry.model_properties_object import ModelProperties, VectorNumericType, Modality, ModelType
+from marqo_commons.model_registry.model_properties_object import ModelProperties, VectorNumericType, Modality, \
+    ModelType, T
 from marqo_commons.model_registry.utils import convert_model_properties_to_dict
 
 
@@ -15,25 +16,26 @@ class FP16ClipModelProperties(ModelProperties):
     modality: List[Modality] = [Modality.image, Modality.text]
     type: ModelType = ModelType.fp16_clip
 
+    @classmethod
+    def get_all_model_properties_objects(cls) -> Dict[str, T]:
+        return {
+            "fp16/ViT-L/14": FP16ClipModelProperties(
+                name="fp16/ViT-L/14",
+                dimensions=768,
+                notes="The faster version (fp16, load from `cuda`) of openai clip model",
+            ),
+            'fp16/ViT-B/32': FP16ClipModelProperties(
+                name="fp16/ViT-B/32",
+                dimensions=512,
+                notes="The faster version (fp16, load from `cuda`) of openai clip model",
+            ),
+            'fp16/ViT-B/16': FP16ClipModelProperties(
+                name="fp16/ViT-B/16",
+                dimensions=512,
+                notes="The faster version (fp16, load from `cuda`) of openai clip model",
+            ),
+        }
 
 @convert_model_properties_to_dict
 def _get_fp16_clip_properties() -> Dict:
-    FP16_CLIP_MODEL_PROPERTIES = {
-        "fp16/ViT-L/14": FP16ClipModelProperties(
-            name="fp16/ViT-L/14",
-            dimensions=768,
-            notes="The faster version (fp16, load from `cuda`) of openai clip model",
-        ),
-        'fp16/ViT-B/32': FP16ClipModelProperties(
-            name="fp16/ViT-B/32",
-            dimensions=512,
-            notes="The faster version (fp16, load from `cuda`) of openai clip model",
-        ),
-        'fp16/ViT-B/16': FP16ClipModelProperties(
-            name="fp16/ViT-B/16",
-            dimensions=512,
-            notes="The faster version (fp16, load from `cuda`) of openai clip model",
-        ),
-    }
-
-    return FP16_CLIP_MODEL_PROPERTIES
+    return FP16ClipModelProperties.get_all_model_properties_objects()