from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.modules import Normalize, Router

query_embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
document_embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

router = Router.for_query_document(
    query_modules=list(query_embedder.children()),
    document_modules=list(document_embedder.children()),
)
model = SentenceTransformer(modules=[router, Normalize()])

query_embedding = model.encode_query(["How do I migrate Asym routes?"])
document_embeddings = model.encode_document(
    [
        "Use Router.for_query_document and encode_query for retrieval queries.",
        "Set the cache directory before loading a local model.",
    ]
)

print("Router routes:", ", ".join(sorted(model[0].sub_modules)))
print("Query embedding shape:", query_embedding.shape)
print("Document embedding shape:", document_embeddings.shape)