#!/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()