Train One Model¶

Use the Python API when you want to train, evaluate, and sample from one model inside a script or notebook.

Notebook¶

Use 01 Run One Model With The Python API for an executable walkthrough in Google Colab.

Data → Model → Fit → Evaluate¶

Three short stages take you from raw splits to scored predictions.

1 Load the splits

from seahorse import load_jsonl

train = load_jsonl("data/my_dataset/train.jsonl")
val   = load_jsonl("data/my_dataset/val.jsonl")
test  = load_jsonl("data/my_dataset/test.jsonl")

2 Build and fit a model

from seahorse import AutoSTPP

model = AutoSTPP(device="cpu", seed=42)
model.fit(train, val, test, epochs=10, batch_size=64, dataset_id="my_dataset")

3 Evaluate and sample

scores  = model.evaluate(test)
samples = model.predict_next(test, n_samples=32)

print(scores)                        # {"test_nll": ..., "mean_seq_nll": ...}
print(samples["next_times"].shape)

Two small rules

fit() requires a validation split. evaluate() returns the Python-facing likelihood metrics test_nll and mean_seq_nll — for benchmark metric profiles and artifact-backed reports, use the CLI.

Variations¶

Compare a baselineSave and reload

from seahorse import PoissonGMM

baseline = PoissonGMM(device="cpu", seed=42)
baseline.fit(train, val, test, epochs=5, batch_size=64)
print(baseline.evaluate(test))

save_dir = model.save("runs/api/auto_stpp")
loaded   = AutoSTPP.load(save_dir)
print(loaded.evaluate(test))

Use Python API for the full Python-facing surface.