from sentence_transformers import SentenceTransformer


def main() -> None:
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    documents = [
        "Reset a forgotten password from the profile security page.",
        "Export paid invoices from the billing dashboard.",
        "Rotate API tokens before sharing an integration.",
        "Change the notification email address for alerts.",
        "Create a vector index for semantic document search.",
        "Archive old support tickets after the retention period.",
        "Review failed background jobs from the worker dashboard.",
        "Update the workspace theme for a user profile.",
    ]
    target_devices = ["cpu", "cpu"]

    pool = model.start_multi_process_pool(target_devices=target_devices)
    try:
        embeddings = model.encode(
            documents,
            pool=pool,
            batch_size=2,
            chunk_size=4,
            normalize_embeddings=True,
            show_progress_bar=False,
        )

        print(f"target devices: {', '.join(target_devices)}")
        print(f"worker processes: {len(pool['processes'])}")
        print(f"documents encoded: {len(documents)}")
        print(f"embedding shape: {embeddings.shape}")
        print(f"first vector norm: {float((embeddings[0] ** 2).sum() ** 0.5):.3f}")
    finally:
        model.stop_multi_process_pool(pool)
        print("pool stopped: True")


if __name__ == "__main__":
    main()