from pathlib import Path

from sentence_transformers import SentenceTransformer


model_id = "sentence-transformers/all-MiniLM-L6-v2"
save_dir = Path("all-minilm-l6-v2-openvino")
openvino_file = "openvino/openvino_model.xml"

model = SentenceTransformer(
    model_id,
    backend="openvino",
    model_kwargs={
        "file_name": openvino_file,
        "export": True,
    },
)

embeddings = model.encode(
    [
        "billing question about an invoice",
        "password reset problem",
    ],
    normalize_embeddings=True,
    show_progress_bar=False,
)

model.save_pretrained(save_dir)

reloaded = SentenceTransformer(
    str(save_dir),
    backend="openvino",
    model_kwargs={
        "file_name": openvino_file,
        "export": False,
    },
    local_files_only=True,
)
reloaded_embedding = reloaded.encode(
    ["billing invoice question"],
    show_progress_bar=False,
)

print(f"backend: {model.get_backend()}")
print(f"embedding shape: {embeddings.shape}")
print(f"saved OpenVINO XML: {save_dir / openvino_file}")
print(f"saved XML exists: {(save_dir / openvino_file).exists()}")
print(f"reloaded backend: {reloaded.get_backend()}")
print(f"reloaded shape: {reloaded_embedding.shape}")