from sentence_transformers import SentenceTransformer, util

documents = [
    {"id": "doc-001", "text": "Set up billing alerts for monthly cloud spending."},
    {"id": "doc-002", "text": "Reset expired password links from the account security page."},
    {"id": "doc-003", "text": "Rotate SSH keys for production deployment hosts."},
    {"id": "doc-004", "text": "Renew TLS certificates before the web server reload."},
    {"id": "doc-005", "text": "Export customer invoices from the finance dashboard."},
]
query = "password reset link expired"

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
corpus = [document["text"] for document in documents]
corpus_embeddings = model.encode_document(
    corpus,
    convert_to_tensor=True,
    normalize_embeddings=True,
    show_progress_bar=False,
)

query_embedding = model.encode_query(
    query,
    convert_to_tensor=True,
    normalize_embeddings=True,
    show_progress_bar=False,
)
hits = util.semantic_search(
    query_embedding,
    corpus_embeddings,
    top_k=3,
)[0]

print(f"Corpus embeddings: {tuple(corpus_embeddings.shape)}")
print(f"Query: {query}")
for rank, hit in enumerate(hits, start=1):
    document = documents[hit["corpus_id"]]
    print(
        f"{rank}. {document['id']} score={hit['score']:.4f} "
        f"text={document['text']}"
    )

top_document = documents[hits[0]["corpus_id"]]
if top_document["id"] != "doc-002":
    raise SystemExit("Semantic search check: failed")
print("Semantic search check: passed")