import os

from sentence_transformers import SentenceTransformer


model_id = os.environ.get("MODEL_ID", "sentence-transformers/all-MiniLM-L6-v2")
model = SentenceTransformer(model_id)

query = "How do I index embeddings for fast semantic search?"
documents = [
    "Use FAISS to build a vector index for dense embeddings.",
    "Use a cross-encoder to rerank a short list of retrieved passages.",
    "Fine-tune an embedding model with hard negatives after collecting labeled pairs.",
]

query_embedding = model.encode_query(
    query,
    normalize_embeddings=True,
    convert_to_tensor=True,
    show_progress_bar=False,
)
document_embeddings = model.encode_document(
    documents,
    normalize_embeddings=True,
    convert_to_tensor=True,
    show_progress_bar=False,
)

scores = model.similarity(query_embedding, document_embeddings)[0]
best_index = int(scores.argmax())

print(f"model_id={model_id}")
print(f"embedding_dimension={model.get_embedding_dimension()}")
print(f"max_sequence_length={model.max_seq_length}")
print(f"query={query}")
print(f"best_match={documents[best_index]}")
print(f"score={float(scores[best_index]):.4f}")

if best_index != 0:
    raise SystemExit(f"unexpected best match index: {best_index}")