Hugging Face Datasets

Datasets from Hugging Face can be converted to PyTRIO Datum objects with the same tokenizer flow used in the training guide.

For example, openai/gsm8k has two configs, main and socratic. Each item contains question and answer fields.

Load the dataset with the datasets library:

from datasets import load_dataset

ds = load_dataset("openai/gsm8k", "main")

train_dataset = ds["train"]
eval_dataset = ds["test"]

Print the first training sample to inspect the schema:

print(train_dataset[0])
"""
{
    'question': 'Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?', 
    'answer': 'Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72'
}
"""

As explained in Datum processing, the core task is to extract text from each dataset item, tokenize it, and package the result as Datum.

Here is an SFT conversion example:

from datasets import load_dataset
import pytrio as trio

train_dataset = load_dataset("openai/gsm8k", "main")["train"]

service_client = trio.ServiceClient()
training_client = service_client.create_lora_training_client(
    base_model="Qwen/Qwen3.5-4B",
    rank=32,
)

print("Loading tokenizer...")
tokenizer = training_client.get_tokenizer()
print("Tokenizer loaded")

def process_example(example: dict, tokenizer) -> trio.Datum:
    prompt = f"Question: {example['question']}\nAnswer:"
    prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True)
    prompt_weights = [0] * len(prompt_tokens)
    
    completion_tokens = tokenizer.encode(f" {example['answer']}\n\n", add_special_tokens=False)
    completion_weights = [1] * len(completion_tokens)
    tokens = prompt_tokens + completion_tokens
    weights = prompt_weights + completion_weights
    input_tokens = tokens[:-1]
    target_tokens = tokens[1:]
    weights = weights[1:]
    
    return trio.Datum(
        model_input=trio.ModelInput.from_ints(tokens=input_tokens),
        loss_fn_inputs=dict(weights=weights, target_tokens=target_tokens)
    )
    
processed_examples = [process_example(ex, tokenizer) for ex in train_dataset]

You can then pass processed_examples to training_client.forward_backward(..., "cross_entropy").