Set cache directory for Hugging Face models in ResNet pipelines.

Author: bojan-karlas

experiments/datascope/experiments/pipelines/pipelines.py

@@ -2,6 +2,7 @@
 import inspect
 import joblib
 import numpy as np
+import os
 import re
 import sklearn.pipeline
 import transformers
@@ -38,6 +39,8 @@
 
 transformers.utils.logging.set_verbosity_error()
 
+HUGGINGFACE_CACHE_DIR = os.path.join("var", "hf-cache")
+
 
 class ProvenancePipeline(sklearn.pipeline.Pipeline):
 
@@ -402,7 +405,7 @@ class ResNet18EmbeddingPipeline(
 
     @classmethod
     def get_model(cls: Type["ResNet18EmbeddingPipeline"]) -> PreTrainedModel:
-        model = ResNetModel.from_pretrained("microsoft/resnet-18")
+        model = ResNetModel.from_pretrained("microsoft/resnet-18", cache_dir=HUGGINGFACE_CACHE_DIR)
         return model
 
 
@@ -415,7 +418,7 @@ class ResNet50EmbeddingPipeline(
 
     @classmethod
     def get_model(cls: Type["ResNet50EmbeddingPipeline"]) -> PreTrainedModel:
-        model = ResNetModel.from_pretrained("microsoft/resnet-50")
+        model = ResNetModel.from_pretrained("microsoft/resnet-50", cache_dir=HUGGINGFACE_CACHE_DIR)
         return model