Skip to content

Commit

Permalink
[cache] If cache_folder nor SENTENCE_TRANSFORMERS_HOME are set,…
Browse files Browse the repository at this point in the history
… use HF default cache (#2412)

* If cache_folder nor SENTENCE_TRANSFORMERS_HOME are set, use HF default cache

This allows models previously loaded in HF (e.g. bert-base-cased) to be automatically used without being downloaded to a second cache.

* Update cache documentation

Although I'm not 100% sure that this still works.
  • Loading branch information
tomaarsen authored Jan 15, 2024
1 parent 3f9ec1e commit e112358
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 13 deletions.
2 changes: 1 addition & 1 deletion docs/training/overview.md
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,7 @@ You can also host the training output on a server and download it:
```python
model = SentenceTransformer('http://www.server.com/path/to/model/my_model.zip')
```
With the first call, the model is downloaded and stored in the local torch cache-folder (`~/.cache/torch/sentence_transformers`). In order to work, you must zip all files and subfolders of your model.
With the first call, the model is downloaded and stored in the local Hugging Face cache folder (`~/.cache/huggingface`). In order to work, you must zip all files and subfolders of your model.



Expand Down
13 changes: 1 addition & 12 deletions sentence_transformers/SentenceTransformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ class SentenceTransformer(nn.Sequential):
SentenceTransformer models from scratch.
:param device: Device (like "cuda", "cpu", "mps") that should be used for computation. If None, checks if a GPU
can be used.
:param cache_folder: Path to store models. Can be also set by SENTENCE_TRANSFORMERS_HOME environment variable.
:param cache_folder: Path to store models. Can also be set by the SENTENCE_TRANSFORMERS_HOME environment variable.
:param trust_remote_code: Whether or not to allow for custom models defined on the Hub in their own modeling files.
This option should only be set to True for repositories you trust and in which you have read the code, as it
will execute code present on the Hub on your local machine.
Expand Down Expand Up @@ -102,17 +102,6 @@ def __init__(

if cache_folder is None:
cache_folder = os.getenv("SENTENCE_TRANSFORMERS_HOME")
if cache_folder is None:
try:
from torch.hub import _get_torch_home

torch_cache_home = _get_torch_home()
except ImportError:
torch_cache_home = os.path.expanduser(
os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"))
)

cache_folder = os.path.join(torch_cache_home, "sentence_transformers")

if model_name_or_path is not None and model_name_or_path != "":
logger.info("Load pretrained SentenceTransformer: {}".format(model_name_or_path))
Expand Down

0 comments on commit e112358

Please sign in to comment.