from pathlib import Path
import importlib as il

import numpy as np


ST = il.import_module(
    "sentence_transformers"
).SentenceTransformer

model_id = (
    "sentence-transformers/"
    "all-MiniLM-L6-v2"
)
save_path = Path(
    "models/"
    "support-embeddings"
)
texts = [
    "reset password",
    "change billing address",
]
files = [
    "modules.json",
    (
        "config_sentence_"
        "transformers.json"
    ),
    "model.safetensors",
    "1_Pooling/config.json",
]

if save_path.exists():
    message = (
        f"{save_path} exists; "
        "choose an empty path"
    )
    raise SystemExit(message)

model = ST(model_id)
before = model.encode(
    texts,
)

model.save_pretrained(
    str(save_path),
    safe_serialization=True,
)

missing = []
for name in files:
    if not (save_path / name).exists():
        missing.append(name)
if missing:
    raise SystemExit(
        "missing saved files: "
        + ", ".join(missing)
    )

reloaded = ST(
    str(save_path),
    local_files_only=True,
)
after = reloaded.encode(
    texts,
)

delta = abs(before - after)
diff = float(
    np.max(delta)
)
if before.shape != after.shape:
    raise SystemExit(
        "shape mismatch"
    )
if diff > 1e-6:
    raise SystemExit(
        "embedding mismatch"
    )

print("saved path: ok")
print(
    "checked files:",
    len(files),
)
print(
    "shape before:",
    before.shape,
)
print(
    "shape after:",
    after.shape,
)
print(
    "max diff:",
    f"{diff:.8f}",
)