from transformers import AutoModel

model_ctx = AutoModel.from_pretrained(
    "seslami-pplx/pplx-embed-context-v1.3.1-4B",
    trust_remote_code=True
)

doc_chunks = [
    [
        "Curiosity begins in childhood with endless questions about the world.",
        "As we grow, curiosity drives us to explore new ideas.",
        "Scientific breakthroughs often start with a curious question."
    ],
    [
        "The curiosity rover explores Mars searching for ancient life.",
        "Each discovery on Mars sparks new questions about the universe."
    ]
]
# Returns list of numpy arrays (one per document)
# embeddings[0].shape = (3, 2560), embeddings[1].shape = (2, 2560)
embeddings = model_ctx.encode(doc_chunks)