Danyil/feature/use marqo commons

requirements.dev.txt

@@ -40,4 +40,5 @@ certifi==2019.11.28
 idna==2.8
 six==1.14.0
 typing-extensions==4.5.0
-urllib3==1.25.8
+urllib3==1.25.8
+marqo-commons @ git+https://github.com/marqo-ai/marqo-commons

src/marqo/s2_inference/model_loaders.py

@@ -0,0 +1,22 @@
+from marqo.s2_inference.hf_utils import HF_MODEL
+from marqo.s2_inference.sbert_onnx_utils import SBERT_ONNX
+from marqo.s2_inference.sbert_utils import SBERT, TEST
+from marqo.s2_inference.random_utils import Random
+from marqo.s2_inference.clip_utils import CLIP, OPEN_CLIP, MULTILINGUAL_CLIP, FP16_CLIP, get_multilingual_clip_properties
+from marqo.s2_inference.types import Any, Dict, List, Optional, Union, FloatTensor
+from marqo.s2_inference.onnx_clip_utils import CLIP_ONNX
+
+# we need to keep track of the embed dim and model load functions/classes
+# we can use this as a registry
+
+def get_model_loaders() -> Dict:
+    return {'clip': CLIP,
+            'open_clip': OPEN_CLIP,
+            'sbert': SBERT,
+            'test': TEST,
+            'sbert_onnx': SBERT_ONNX,
+            'clip_onnx': CLIP_ONNX,
+            "multilingual_clip": MULTILINGUAL_CLIP,
+            "fp16_clip": FP16_CLIP,
+            'random': Random,
+            'hf': HF_MODEL}

src/marqo/s2_inference/onnx_clip_utils.py

@@ -14,7 +14,6 @@
 from marqo.s2_inference.logger import get_logger
 import onnxruntime as ort
 from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-import marqo.s2_inference.model_registry as model_registry
 from zipfile import ZipFile
 from huggingface_hub.utils import RevisionNotFoundError,RepositoryNotFoundError, EntryNotFoundError, LocalEntryNotFoundError
 from marqo.s2_inference.errors import ModelDownloadError
@@ -60,6 +59,7 @@ class CLIP_ONNX(object):
     def __init__(self, model_name: str ="onnx32/openai/ViT-L/14", device: str = None, embedding_dim: int = None,
                  truncate: bool = True,
                  load=True, **kwargs):
+        from marqo.s2_inference.s2_inference import get_model_properties_from_registry
         self.model_name = model_name
         self.onnx_type, self.source, self.clip_model = self.model_name.split("/", 2)
         if not device:
@@ -70,7 +70,7 @@ def __init__(self, model_name: str ="onnx32/openai/ViT-L/14", device: str = None
             "CPUExecutionProvider"]
         self.visual_session = None
         self.textual_session = None
-        self.model_info = model_registry._get_onnx_clip_properties()[self.model_name]
+        self.model_info = get_model_properties_from_registry(self.model_name)
 
         self.visual_type = np.float16 if self.onnx_type == "onnx16" else np.float32
         self.textual_type = np.int64 if self.source == "open_clip" else np.int32

src/marqo/s2_inference/s2_inference.py

@@ -2,12 +2,14 @@
 The functions defined here would have endpoints, later on.
 """
 import numpy as np
+from marqo_commons.model_registry.model_registry import get_model_properties_dict
+
 from marqo.errors import ModelCacheManagementError, InvalidArgError, ConfigurationError, InternalError
 from marqo.s2_inference.errors import (
     VectoriseError, InvalidModelPropertiesError, ModelLoadError,
     UnknownModelError, ModelNotInCacheError, ModelDownloadError, S2InferenceError)
 from PIL import UnidentifiedImageError
-from marqo.s2_inference.model_registry import load_model_properties
+from marqo.s2_inference.model_loaders import get_model_loaders
 from marqo.s2_inference.configs import get_default_normalization, get_default_seq_length
 from marqo.s2_inference.types import *
 from marqo.s2_inference.logger import get_logger
@@ -28,7 +30,8 @@
 available_models = dict()
 # A lock to protect the model loading process
 lock = threading.Lock()
-MODEL_PROPERTIES = load_model_properties()
+MODEL_PROPERTIES = get_model_properties_dict()
+MODEL_LOADERS = get_model_loaders()
 
 
 def vectorise(model_name: str, content: Union[str, List[str]], model_properties: dict = None,
@@ -300,8 +303,8 @@ def get_model_size(model_name: str, model_properties: dict) -> (int, float):
     Return the model size for given model
     Note that the priorities are size_in_properties -> model_name -> model_type -> default size
     '''
-    if "model_size" in model_properties:
-        return model_properties["model_size"]
+    if "memory_size" in model_properties:
+        return model_properties["memory_size"]
 
     name_info = (model_name + model_properties.get("name", "")).lower().replace("/", "-")
     for name, size in constants.MODEL_NAME_SIZE_MAPPING.items():
@@ -369,11 +372,11 @@ def get_model_properties_from_registry(model_name: str) -> dict:
         dict: a dictionary describing properties of the model.
     """
 
-    if model_name not in MODEL_PROPERTIES['models']:
+    if model_name not in MODEL_PROPERTIES:
         raise UnknownModelError(f"Could not find model properties in model registry for model={model_name}. "
                                 f"Model is not supported by default.")
 
-    return MODEL_PROPERTIES['models'][model_name]
+    return MODEL_PROPERTIES[model_name]
 
 
 def _check_output_type(output: List[List[float]]) -> bool:
@@ -505,10 +508,10 @@ def _get_model_loader(model_name: str, model_properties: dict) -> Any:
 
     model_type = model_properties['type']
 
-    if model_type not in MODEL_PROPERTIES['loaders']:
+    if model_type not in MODEL_LOADERS:
         raise KeyError(f"model_name={model_name} for model_type={model_type} not in allowed model types")
 
-    return MODEL_PROPERTIES['loaders'][model_type]
+    return MODEL_LOADERS[model_type]
 
 
 def get_available_models():