from sentence_transformers import SentenceTransformer
import numpy as np


model = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2",
    prompts={
        "query": "query: ",
        "document": "passage: ",
    },
)

query = "How do I reset a forgotten password?"
documents = [
    "Generate quarterly revenue charts from a CSV export.",
    "Reset a lost account password from the profile security page.",
    "Tune the database connection pool for a busy API server.",
]

query_embedding = model.encode_query(
    [query],
    normalize_embeddings=True,
    show_progress_bar=False,
)
document_embeddings = model.encode_document(
    documents,
    normalize_embeddings=True,
    show_progress_bar=False,
)
inline_query_embedding = model.encode(
    [query],
    prompt="query: ",
    normalize_embeddings=True,
    show_progress_bar=False,
)

scores = model.similarity(query_embedding, document_embeddings)[0]
best_index = int(scores.argmax())
inline_delta = float(np.max(np.abs(query_embedding - inline_query_embedding)))

print(f"prompt keys: {', '.join(sorted(model.prompts))}")
print(f"query shape: {query_embedding.shape}")
print(f"document shape: {document_embeddings.shape}")
print(f"inline prompt delta: {inline_delta:.6f}")
print(f"top match: doc-{best_index + 1}")
print(f"score: {scores[best_index]:.3f}")
print(f"text: {documents[best_index]}")