API
pytrio.Datum
class Datum(BaseModel):
model_input: ModelInput
loss_fn_inputs: dict[str, TensorData]Datum is the data structure for one training sample. Pass a list of Datum objects to methods such as TrainingClient.forward_backward().
Each Datum contains input token IDs and the parameters required by the selected loss function.
Fields
| Field | Type | Description |
|---|---|---|
model_input | ModelInput | Input token ID list |
loss_fn_inputs | dict[str, TensorData | Any] | Loss-function inputs. Keys depend on the selected loss function |
loss_fn_inputs
Different loss functions require different keys in loss_fn_inputs.
cross_entropy (default)
| Key | Type | Required | Description |
|---|---|---|---|
target_tokens | TensorData | Yes | Target token IDs. Length must match model_input |
weights | TensorData | No | Per-position loss weights. Defaults to all 1.0 |
importance_sampling / ppo
| Key | Type | Required | Description |
|---|---|---|---|
target_tokens | TensorData | Yes | Target token IDs |
logprobs | TensorData | Yes | Logprobs from the old policy or sampler |
advantages | TensorData | Yes | Advantage values |
All three fields should have the same length as model_input.
Examples
cross_entropy
import pytrio as trio
tokens = tokenizer.encode("The meaning of life is")
input_tokens = tokens[:-1]
target_tokens = tokens[1:]
datum = trio.Datum(
model_input=trio.ModelInput.from_ints(input_tokens),
loss_fn_inputs={"target_tokens": target_tokens},
)weighted cross_entropy
datum = trio.Datum(
model_input=trio.ModelInput.from_ints(input_tokens),
loss_fn_inputs={
"target_tokens": target_tokens,
"weights": [0.0] * 5 + [1.0] * (len(target_tokens) - 5), # Exclude the first 5 tokens from the loss
},
)importance_sampling / ppo
datum = trio.Datum(
model_input=trio.ModelInput.from_ints(input_tokens),
loss_fn_inputs={
"target_tokens": target_tokens,
"logprobs": old_logprobs,
"advantages": advantages,
},
)Was this documentation helpful?