Integrate with Sentence Transformers v5.4

sentence_transformers_impl.py

@@ -54,17 +54,21 @@ class Transformer(nn.Module):
         if config_args is None:
             config_args = {}
 
+        if cache_dir is not None:
+            config_args["cache_dir"] = cache_dir
+            model_args["cache_dir"] = cache_dir
+            tokenizer_args["cache_dir"] = cache_dir
+
         if not model_args.get("trust_remote_code", False):
             raise ValueError(
                 "You need to set `trust_remote_code=True` to load this model."
             )
 
-        self.config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir)
-        self.auto_model = AutoModel.from_pretrained(model_name_or_path, config=self.config, cache_dir=cache_dir, **model_args)
+        self.config = AutoConfig.from_pretrained(model_name_or_path, **config_args)
+        self.auto_model = AutoModel.from_pretrained(model_name_or_path, config=self.config, **model_args)
 
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_name_or_path,
-            cache_dir=cache_dir,
             **tokenizer_args,
         )