dataset

Real Puzzle-KD text dataset adapter for native AutoModel global KD.

Functions

collate_puzzletron_llm_batch

Stack fixed-length tensor samples without list-oriented padding.

make_puzzletron_chat_dataset

Load saved, Hub, JSON, or Parquet messages through AutoModel chat formatting.

make_puzzletron_llm_dataset

Pack the real local Puzzle-KD messages into deterministic fixed-length samples.

make_puzzletron_llm_overfit_dataset

Materialize one deterministic minibatch and replay it every epoch.
collate_puzzletron_llm_batch(batch)

Stack fixed-length tensor samples without list-oriented padding.

Puzzletron’s packed-token cache already emits equally sized tensors. The generic AutoModel collator is intended for variable-length Python lists and attempts tensor + list while padding these samples. Keep this collator dataset-specific and preserve AutoModel’s usual padding_mask contract.

Parameters:

batch (list[dict[str, Tensor]])

Return type:

dict[str, Tensor]

make_puzzletron_chat_dataset(tokenizer, dataset_path, split='train', num_samples=None, seq_length=4096, seed=444, **_)

Load saved, Hub, JSON, or Parquet messages through AutoModel chat formatting.

Parameters:

  • dataset_path (str)

  • split (str)

  • num_samples (int | None)

  • seq_length (int)

  • seed (int)

  • _ (object)

make_puzzletron_llm_dataset(tokenizer, dataset_path, split='train', num_samples=2048, seq_length=131072, seed=444, packed_token_cache_path=None, **_)

Pack the real local Puzzle-KD messages into deterministic fixed-length samples.

Parameters:

  • dataset_path (str)

  • split (str)

  • num_samples (int)

  • seq_length (int)

  • seed (int)

  • packed_token_cache_path (str | None)

  • _ (object)

Return type:

IterableDataset

make_puzzletron_llm_overfit_dataset(tokenizer, dataset_path, split='train', num_samples=128, seq_length=128, seed=444, packed_token_cache_path=None, **kwargs)

Materialize one deterministic minibatch and replay it every epoch.

Parameters:

  • dataset_path (str)

  • split (str)

  • num_samples (int)

  • seq_length (int)

  • seed (int)

  • packed_token_cache_path (str | None)

  • kwargs (object)

Return type:

IterableDataset