blob: 6478a507a9750a3849efa3f45fefc1aa34814bff (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
|
#!/usr/bin/env python3
from __future__ import annotations
import os
from pathlib import Path
from sentence_transformers import SentenceTransformer
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
LOCAL_CACHE_DIR = Path("models/hf")
def main() -> None:
LOCAL_CACHE_DIR.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HF_HOME", str(LOCAL_CACHE_DIR.resolve()))
os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", str(LOCAL_CACHE_DIR.resolve()))
print(f"Caching model: {MODEL_NAME}")
print(f"Cache dir: {LOCAL_CACHE_DIR.resolve()}")
model = SentenceTransformer(
MODEL_NAME,
cache_folder=str(LOCAL_CACHE_DIR.resolve()),
)
# Force an actual encode call so all needed files are loaded.
_ = model.encode(["test"], convert_to_numpy=True)
print("Model cached successfully.")
if __name__ == "__main__":
main()
|
