Rad Dino changes

Azure · Dec 9, 2024 · 3f79d30 · 3f79d30
1 parent 317261a
commit 3f79d30
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 1 deletion.
diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -79,6 +79,7 @@ inputs:
       - mask-generation
       - video-multi-object-tracking
       - visual-question-answering
+      - embeddings
     description: A Hugging face task on which model was trained on. A required parameter for transformers MLflow flavor. Can be provided as input here or in model_download_metadata JSON file.
     optional: true
 

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
@@ -100,6 +100,7 @@ inputs:
       - mask-generation
       - video-multi-object-tracking
       - visual-question-answering
+      - embeddings
     optional: true
     type: string
 

diff --git a/assets/training/model_management/src/azureml/model/mgmt/processors/factory.py b/assets/training/model_management/src/azureml/model/mgmt/processors/factory.py
@@ -72,7 +72,7 @@ def get_mlflow_convertor(model_framework, model_dir, output_dir, temp_dir, trans
             return CLIPMLflowConvertorFactory.create_mlflow_convertor(
                 model_dir, output_dir, temp_dir, translate_params
             )
-        elif (task == PyFuncSupportedTasks.EMBEDDINGS.value) and model_id.startswith(ModelFamilyPrefixes.DINOV2.value):
+        elif (task == PyFuncSupportedTasks.EMBEDDINGS.value) and (model_id.startswith(ModelFamilyPrefixes.DINOV2.value) or model_id.startswith(ModelFamilyPrefixes.RADDINO.value)):
             return DinoV2MLflowConvertorFactory.create_mlflow_convertor(
                 model_dir, output_dir, temp_dir, translate_params
             )

diff --git a/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/config.py b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/config.py
@@ -65,3 +65,5 @@ class ModelFamilyPrefixes(_CustomEnum):
 
     # DinoV2 model family.
     DINOV2 = "facebook/dinov2"
+
+    RADDINO = "microsoft/rad-dino"
Original file line number	Diff line number	Diff line change
Expand Up		@@ -65,3 +65,5 @@ class ModelFamilyPrefixes(_CustomEnum):

		# DinoV2 model family.
		DINOV2 = "facebook/dinov2"

		RADDINO = "microsoft/rad-dino"