File size: 1,810 Bytes
4e70c09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | from __future__ import annotations
from typing import Any, Protocol
import torch
from torch.utils.data import Dataset
from conveyer.config import ConveyerConfig
class TokenizerProtocol(Protocol):
vocab_size: int
pad_id: int
bos_id: int
eos_id: int
def encode(self, text: str) -> list[int]: ...
def decode(self, ids: list[int]) -> str: ...
def id_to_piece(self, token_id: int) -> str: ...
def build_dataset_and_tokenizer(
cfg: ConveyerConfig,
infer_only: bool = False,
) -> tuple[Dataset[dict[str, Any]], Dataset[dict[str, Any]], TokenizerProtocol]:
if cfg.data.dataset == "synthetic":
from conveyer.data.synthetic import build_synthetic_datasets
return build_synthetic_datasets(cfg)
elif cfg.data.dataset == "tinystories":
from conveyer.data.tinystories import build_tinystories_datasets
return build_tinystories_datasets(cfg, infer_only=infer_only)
else:
raise ValueError(f"Unknown dataset '{cfg.data.dataset}'. Choose 'synthetic' or 'tinystories'.")
def collate_conveyer_batch(batch: list[dict[str, Any]], pad_id: int = 0) -> dict[str, torch.Tensor]:
max_len = max(item["input_ids"].size(0) for item in batch)
bsz = len(batch)
input_ids = torch.full((bsz, max_len), pad_id, dtype=torch.long)
prompt_lengths = torch.zeros(bsz, dtype=torch.long)
task_ids = torch.zeros(bsz, dtype=torch.long)
for i, item in enumerate(batch):
seq = item["input_ids"]
input_ids[i, : seq.size(0)] = seq
prompt_lengths[i] = int(item.get("prompt_len", 1))
task_ids[i] = int(item.get("task_id", 0))
return {
"input_ids": input_ids,
"prompt_lengths": prompt_lengths,
"task_ids": task_ids,
}
|