from sentence_transformers import SparseEncoder


model = SparseEncoder("rasyosef/splade-tiny", max_active_dims=64)

documents = [
    "Reset expired password links from the account security page.",
    "Renew TLS certificates before the web server reload.",
    "Export customer invoices from the finance dashboard.",
]
query = "web server certificate renewal"

document_embeddings = model.encode_document(
    documents,
    convert_to_sparse_tensor=True,
    show_progress_bar=False,
)
query_embedding = model.encode_query(
    [query],
    convert_to_sparse_tensor=True,
    show_progress_bar=False,
)

document_stats = SparseEncoder.sparsity(document_embeddings)
query_stats = SparseEncoder.sparsity(query_embedding)
query_tokens = model.decode(query_embedding, top_k=4)[0]
scores = model.similarity(query_embedding, document_embeddings)[0]
best_index = int(scores.argmax())

print(f"document shape: {tuple(document_embeddings.shape)}")
print(f"query shape: {tuple(query_embedding.shape)}")
print(f"document active dims: {document_stats['active_dims']:.1f}")
print(f"query active dims: {query_stats['active_dims']:.1f}")
print(f"query sparsity: {query_stats['sparsity_ratio']:.4f}")
print("top query tokens:")
for token, weight in query_tokens:
    print(f"  {token}: {weight:.3f}")
print(f"top match: doc-{best_index + 1}")
print(f"text: {documents[best_index]}")