Repository: microsoft/hi-ml

Issue:
Specify commit hash when downloading multimodal model weights
A warning is shown when downloading models from Hugging Face:

```text
Explicitly passing a `revision` is encouraged when loading a configuration with custom code to ensure no malicious code has been contributed in a newer revision.
Explicitly passing a `revision` is encouraged when loading a model with custom code to ensure no malicious code has been contributed in a newer revision.
```

We need to specify the `revision` for security.


Discussion:

Resolving patch:
diff --git a/hi-ml-multimodal/src/health_multimodal/text/__init__.py b/hi-ml-multimodal/src/health_multimodal/text/__init__.py
--- a/hi-ml-multimodal/src/health_multimodal/text/__init__.py
+++ b/hi-ml-multimodal/src/health_multimodal/text/__init__.py
@@ -11,6 +11,7 @@
    :toctree:
 
    inference_engine
+   utils
 
 
 .. currentmodule:: health_multimodal.text.data
@@ -28,14 +29,25 @@
 
    configuration_cxrbert
    modelling_cxrbert
+
 """
 
 from .data.io import TypePrompts
+from .inference_engine import TextInferenceEngine
+from .model.modelling_cxrbert import CXRBertModel, CXRBertOutput
+from .model.configuration_cxrbert import CXRBertConfig, CXRBertTokenizer
 
 
 BIOMED_VLP_CXR_BERT_SPECIALIZED = "microsoft/BiomedVLP-CXR-BERT-specialized"
+CXR_BERT_COMMIT_TAG = "v1.0"
 
 __all__ = [
     "BIOMED_VLP_CXR_BERT_SPECIALIZED",
+    "CXR_BERT_COMMIT_TAG",
     "TypePrompts",
+    "TextInferenceEngine",
+    "CXRBertConfig",
+    "CXRBertTokenizer",
+    "CXRBertModel",
+    "CXRBertOutput",
 ]
diff --git a/hi-ml-multimodal/src/health_multimodal/text/utils.py b/hi-ml-multimodal/src/health_multimodal/text/utils.py
new file mode 100644
--- /dev/null
+++ b/hi-ml-multimodal/src/health_multimodal/text/utils.py
@@ -0,0 +1,35 @@
+#  -------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+#  -------------------------------------------------------------------------------------------
+
+
+from typing import Tuple
+
+from health_multimodal.text import (
+    BIOMED_VLP_CXR_BERT_SPECIALIZED,
+    CXR_BERT_COMMIT_TAG,
+    TextInferenceEngine,
+    CXRBertModel,
+    CXRBertTokenizer,
+)
+
+
+def get_cxr_bert() -> Tuple[CXRBertTokenizer, CXRBertModel]:
+    """Load the CXR-BERT model and tokenizer from the `Hugging Face Hub <https://huggingface.co/microsoft/BiomedVLP-CXR-BERT-specialized>`_."""  # noqa: E501
+    model_name = BIOMED_VLP_CXR_BERT_SPECIALIZED
+    revision = CXR_BERT_COMMIT_TAG
+    tokenizer = CXRBertTokenizer.from_pretrained(model_name, revision=revision)
+    text_model = CXRBertModel.from_pretrained(model_name, revision=revision)
+    return tokenizer, text_model
+
+
+def get_cxr_bert_inference() -> TextInferenceEngine:
+    """Create a :class:`TextInferenceEngine` for the CXR-BERT model.
+
+    The model is downloaded from the Hugging Face Hub.
+    The engine can be used to get embeddings from text prompts or masked token predictions.
+    """
+    tokenizer, text_model = get_cxr_bert()
+    text_inference = TextInferenceEngine(tokenizer=tokenizer, text_model=text_model)
+    return text_inference


Test patch:
diff --git a/hi-ml-multimodal/test_multimodal/text/data/test_text_io.py b/hi-ml-multimodal/test_multimodal/text/data/test_text_io.py
--- a/hi-ml-multimodal/test_multimodal/text/data/test_text_io.py
+++ b/hi-ml-multimodal/test_multimodal/text/data/test_text_io.py
@@ -4,16 +4,11 @@
 #  -------------------------------------------------------------------------------------------
 
 import pytest
-from transformers import AutoModel, AutoTokenizer
 
-from health_multimodal.text import BIOMED_VLP_CXR_BERT_SPECIALIZED, TypePrompts
-from health_multimodal.text.inference_engine import TextInferenceEngine
+from health_multimodal.text import TypePrompts
+from health_multimodal.text.utils import get_cxr_bert_inference
 
-
-text_inference = TextInferenceEngine(
-    tokenizer=AutoTokenizer.from_pretrained(BIOMED_VLP_CXR_BERT_SPECIALIZED, trust_remote_code=True),
-    text_model=AutoModel.from_pretrained(BIOMED_VLP_CXR_BERT_SPECIALIZED, trust_remote_code=True),
-)
+text_inference = get_cxr_bert_inference()
 
 
 @pytest.mark.parametrize("prompts", ("", "hello", "this is a test", ["this is", "also a test"]))
diff --git a/hi-ml-multimodal/test_multimodal/text/test_inference_engine.py b/hi-ml-multimodal/test_multimodal/text/test_inference_engine.py
--- a/hi-ml-multimodal/test_multimodal/text/test_inference_engine.py
+++ b/hi-ml-multimodal/test_multimodal/text/test_inference_engine.py
@@ -4,22 +4,18 @@
 #  -------------------------------------------------------------------------------------------
 
 
-from typing import Tuple
-
 import pytest
 import torch
 
-from health_multimodal.text import BIOMED_VLP_CXR_BERT_SPECIALIZED
 from health_multimodal.text.inference_engine import TextInferenceEngine
-from health_multimodal.text.model.modelling_cxrbert import CXRBertModel
-from health_multimodal.text.model.configuration_cxrbert import CXRBertTokenizer
+from health_multimodal.text.utils import get_cxr_bert
 
 
 def test_text_inference_init_model_type() -> None:
     """
     Test that init fails if the wrong model type is passed in
     """
-    tokenizer, _ = _get_cxr_bert()
+    tokenizer, _ = get_cxr_bert()
     false_model = torch.nn.Linear(4, 4)
     with pytest.raises(AssertionError) as ex:
         TextInferenceEngine(tokenizer=tokenizer, text_model=false_model)  # type: ignore[arg-type]
@@ -30,7 +26,7 @@ def test_l2_normalization() -> None:
     """
     Test that the text embeddings (CLS token) are l2 normalized.
     """
-    tokenizer, text_model = _get_cxr_bert()
+    tokenizer, text_model = get_cxr_bert()
 
     text_inference = TextInferenceEngine(tokenizer=tokenizer, text_model=text_model)
     input_query = ["There is a tumor in the left lung", "Lungs are all clear"]
@@ -43,7 +39,7 @@ def test_sentence_semantic_similarity() -> None:
     """
     Test that the sentence embedding similarity computed by the text model is meaningful.
     """
-    tokenizer, text_model = _get_cxr_bert()
+    tokenizer, text_model = get_cxr_bert()
 
     # CLS token has no dedicated meaning, but we can expect vector similarity due to token overlap between the sentences
     text_inference = TextInferenceEngine(tokenizer=tokenizer, text_model=text_model)
@@ -56,19 +52,12 @@ def test_sentence_semantic_similarity() -> None:
     assert pos_sim > neg_sim_2
 
 
-def _get_cxr_bert() -> Tuple[CXRBertTokenizer, CXRBertModel]:
-    model_name = BIOMED_VLP_CXR_BERT_SPECIALIZED
-    tokenizer = CXRBertTokenizer.from_pretrained(model_name)
-    text_model = CXRBertModel.from_pretrained(model_name)
-    return tokenizer, text_model
-
-
 def test_triplet_similarities_with_inference_engine() -> None:
     """
     Test that the triplet sentence similarities computed by the text model are meaningful.
     """
 
-    tokenizer, text_model = _get_cxr_bert()
+    tokenizer, text_model = get_cxr_bert()
     text_inference = TextInferenceEngine(tokenizer=tokenizer, text_model=text_model)
 
     reference = \
@@ -102,7 +91,7 @@ def test_mlm_with_inference_engine_with_hf_hub() -> None:
     Test that the MLM model can be used with the inference engine and the filled masked tokens are correct.
     """
 
-    tokenizer, text_model = _get_cxr_bert()
+    tokenizer, text_model = get_cxr_bert()
     text_inference = TextInferenceEngine(tokenizer=tokenizer, text_model=text_model)
 
     # ##### Test Masked Language Modelling ######
diff --git a/hi-ml-multimodal/test_multimodal/vlp/test_vlp_inference_engine.py b/hi-ml-multimodal/test_multimodal/vlp/test_vlp_inference_engine.py
--- a/hi-ml-multimodal/test_multimodal/vlp/test_vlp_inference_engine.py
+++ b/hi-ml-multimodal/test_multimodal/vlp/test_vlp_inference_engine.py
@@ -10,21 +10,13 @@
 import pytest
 import numpy as np
 from PIL import Image
-from transformers import AutoModel, AutoTokenizer
 
-from health_multimodal.text import BIOMED_VLP_CXR_BERT_SPECIALIZED
-from health_multimodal.text.inference_engine import TextInferenceEngine
+from health_multimodal.text.utils import get_cxr_bert_inference
 from health_multimodal.image import ImageModel, ResnetType, ImageInferenceEngine
 from health_multimodal.image.data.transforms import create_chest_xray_transform_for_inference
 from health_multimodal.vlp.inference_engine import ImageTextInferenceEngine
 
 
-text_inference = TextInferenceEngine(
-    tokenizer=AutoTokenizer.from_pretrained(BIOMED_VLP_CXR_BERT_SPECIALIZED, trust_remote_code=True),
-    text_model=AutoModel.from_pretrained(BIOMED_VLP_CXR_BERT_SPECIALIZED, trust_remote_code=True),
-)
-
-
 @pytest.mark.parametrize("height", (400, 500, 650))
 @pytest.mark.parametrize("query_text", ("", "hello", "this is a test"))
 def test_vlp_inference(height: int, query_text: str) -> None:
@@ -42,7 +34,7 @@ def test_vlp_inference(height: int, query_text: str) -> None:
         transform=create_chest_xray_transform_for_inference(resize=resize, center_crop_size=center_crop_size))
     img_txt_inference = ImageTextInferenceEngine(
         image_inference_engine=image_inference,
-        text_inference_engine=text_inference,
+        text_inference_engine=get_cxr_bert_inference(),
     )
 
     with tempfile.NamedTemporaryFile(suffix='.jpg') as f: