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

assert query_embedding.shape[1] == document_embeddings.shape[1]
assert best_index == 1

print(f"query shape: {query_embedding.shape}")
print(f"document shape: {document_embeddings.shape}")
print(f"top match: doc-{best_index + 1}")
print(f"score: {scores[best_index]:.3f}")
print(f"text: {documents[best_index]}")