import sentence_transformers
from sentence_transformers import SentenceTransformer

modules = sentence_transformers.sentence_transformer.modules
Transformer = modules.Transformer
Pooling = modules.Pooling

MODEL_NAME = (
    "sentence-transformers/"
    "all-MiniLM-L6-v2"
)

model = SentenceTransformer(
    MODEL_NAME,
    processor_kwargs={"model_max_length": 128},
)

query_embedding = model.encode_query(
    ["How do I migrate code?"]
)
document_embeddings = model.encode_document(
    [
        (
            "Use processor_kwargs "
            "instead of tokenizer_kwargs."
        ),
        "Use encode_query.",
    ]
)
truncated_embeddings = model.encode(
    inputs=[
        (
            "Sentence Transformers "
            "v5 migration check"
        )
    ],
    truncate_dim=128,
)

print("Dimension:", model.get_embedding_dimension())
print("Max sequence length:", model.max_seq_length)
print("Query shape:", query_embedding.shape)
print("Document shape:", document_embeddings.shape)
print("Truncated shape:", truncated_embeddings.shape)
print("Imported modules:", Transformer.__name__, Pooling.__name__)