fix: kwargs in custom Sentence Transformer

custom_st.py

@@ -22,6 +22,8 @@ class Transformer(nn.Module):
         model_kwargs: Optional[Dict[str, Any]] = None,
         tokenizer_kwargs: Optional[Dict[str, Any]] = None,
         image_processor_kwargs: Optional[Dict[str, Any]] = None,
+        cache_dir: str = None,
+        **_,
     ) -> None:
         super(Transformer, self).__init__()
 
@@ -30,19 +32,23 @@ class Transformer(nn.Module):
         tokenizer_kwargs = tokenizer_kwargs or {}
         image_processor_kwargs = image_processor_kwargs or {}
 
-        config = AutoConfig.from_pretrained(model_name_or_path, **config_kwargs)
+        config = AutoConfig.from_pretrained(
+            model_name_or_path, cache_dir=cache_dir, **config_kwargs
+        )
         self.model = AutoModel.from_pretrained(
-            model_name_or_path, config=config, **model_kwargs
+            model_name_or_path, config=config, cache_dir=cache_dir, **model_kwargs
         )
         if max_seq_length is not None and 'model_max_length' not in tokenizer_kwargs:
             tokenizer_kwargs['model_max_length'] = max_seq_length
 
         self.tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_name_or_path or model_name_or_path,
+            cache_dir=cache_dir,
             **tokenizer_kwargs,
         )
         self.image_processor = AutoImageProcessor.from_pretrained(
             image_processor_name_or_path or model_name_or_path,
+            cache_dir=cache_dir,
             **image_processor_kwargs,
         )