train_result = trainer.train()
model.save_pretrained(model_dir)

trained_model = SparseEncoder(str(model_dir))
query_embedding = trained_model.encode_query(
    ["How do I replace an expired API token?"],
    convert_to_tensor=True,
).coalesce()
document_embeddings = trained_model.encode_document(
    [
        "Revoke the old API token and create a replacement token.",
        "Open billing settings and update the saved credit card.",
        "Open team settings and send an invitation email.",
    ],
    convert_to_tensor=True,
).coalesce()

scores = torch.mm(
    query_embedding.to_dense(),
    document_embeddings.to_dense().T,
)
best_index = int(scores.argmax().item())

print(f"Training rows: {len(train_dataset)}")
print(f"Loss wrapper: {loss.__class__.__name__}")
print(f"Training steps: {train_result.global_step}")
print(f"Saved model: {model_dir}")
print(f"Query active dimensions: {query_embedding._nnz()}")
print(f"Document active dimensions: {document_embeddings._nnz()}")
print(f"Top smoke-test document: d{best_index + 1}")
print(f"Top smoke-test score: {scores[0, best_index]:.4f}")

if train_result.global_step != 2:
    raise SystemExit(f"unexpected training steps: {train_result.global_step}")
if query_embedding._nnz() == 0 or document_embeddings._nnz() == 0:
    raise SystemExit("sparse encoder returned no active dimensions")
if not torch.isfinite(scores).all():
    raise SystemExit("non-finite sparse similarity score returned")
if best_index != 0:
    raise SystemExit("the API token query did not rank the matching document first")