from transformers import AutoModel, AutoTokenizer

REPO_ID = "YOUR_ORG/LaGrange-1.0-Flash"
REVISION = "main"

tokenizer = AutoTokenizer.from_pretrained(
    REPO_ID,
    revision=REVISION,
    trust_remote_code=True,
)
model = AutoModel.from_pretrained(
    REPO_ID,
    revision=REVISION,
    trust_remote_code=True,
)

batch = tokenizer("machine learning mathematics", return_tensors="pt")
bridge_output = model(**batch)
print(bridge_output.pooler_output.shape)

result = model.research({
    "domain": "ml_math",
    "action": "softmax_jacobian",
    "logits": [0.2, -0.1, 1.3, 2.0],
})
print(result)