from pathlib import Path

import onnxruntime as ort
from sentence_transformers import SentenceTransformer


model_id = "sentence-transformers/all-MiniLM-L6-v2"
provider = "CPUExecutionProvider"
save_dir = Path("all-minilm-l6-v2-onnx")

available_providers = ort.get_available_providers()
if provider not in available_providers:
    raise SystemExit(f"{provider} is not available")

model = SentenceTransformer(
    model_id,
    backend="onnx",
    model_kwargs={
        "provider": provider,
        "file_name": "onnx/model.onnx",
    },
)
embeddings = model.encode(
    [
        "billing question about an invoice",
        "password reset problem",
    ],
    normalize_embeddings=True,
    show_progress_bar=False,
)

print(f"backend: {model.get_backend()}")
print(f"available providers: {', '.join(available_providers)}")
print(f"embedding shape: {embeddings.shape}")

model.save_pretrained(save_dir)
print(f"saved ONNX file: {save_dir / 'onnx' / 'model.onnx'}")

reloaded = SentenceTransformer(
    str(save_dir),
    backend="onnx",
    model_kwargs={
        "provider": provider,
        "file_name": "onnx/model.onnx",
    },
    local_files_only=True,
)
reloaded_embedding = reloaded.encode(
    ["billing invoice question"],
    show_progress_bar=False,
)

print(f"reloaded backend: {reloaded.get_backend()}")
print(f"reloaded shape: {reloaded_embedding.shape}")