from sentence_transformers import CrossEncoder, SentenceTransformer, util

query = "How do I restore a Docker volume backup on another host?"

documents = [
    "Create a tar archive from a Docker volume, copy it to the new host, and extract it into a replacement volume.",
    "Use docker compose pull and docker compose up -d to recreate application containers after an image update.",
    "List Docker images and tags before promoting a release to production.",
    "Upload a finished build directory to Amazon S3 with aws s3 sync for release backups.",
    "Inspect container logs with docker logs when a service exits during startup.",
]

retriever = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2", device="cpu")

document_embeddings = retriever.encode(
    documents,
    normalize_embeddings=True,
    show_progress_bar=False,
)
query_embedding = retriever.encode(
    query,
    normalize_embeddings=True,
    show_progress_bar=False,
)

hits = util.semantic_search(query_embedding, document_embeddings, top_k=4)[0]
candidates = [documents[hit["corpus_id"]] for hit in hits]

print("Initial embedding search:")
for position, hit in enumerate(hits, start=1):
    print(f"{position}. cosine={hit['score']:.3f} | {documents[hit['corpus_id']]}")

print()
print("Final reranked order:")
for position, hit in enumerate(
    reranker.rank(
        query,
        candidates,
        return_documents=True,
        show_progress_bar=False,
    ),
    start=1,
):
    print(f"{position}. score={hit['score']:.2f} | {hit['text']}")