File size: 7,841 Bytes
31dc8dc | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | import json
from typing import Dict, Optional
import jax
import jax.numpy as jnp
import numpy as np
from datasets import DatasetDict, load_dataset as hf_load_dataset, load_from_disk
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from utils.encoder_utils import build_self_attn_cond_masks
from utils.logging_utils import log_for_0
PRNGKey = jax.random.PRNGKey
def get_pad_token_id(tokenizer, pad_token="pad"):
"""Resolve the token id used for padding, optionally using EOS as pad."""
token_id = tokenizer.eos_token_id if pad_token == "eos" else tokenizer.pad_token_id
if token_id is None:
raise ValueError(
f"Tokenizer has no pad_token_id or eos_token_id."
)
return token_id
def prepare_batch(batch: Dict, config, rng: PRNGKey):
"""Convert numpy batch to JAX arrays and sample label-drop decisions."""
result = {k: jnp.array(v) if isinstance(v, np.ndarray) else v for k, v in batch.items()}
input_ids = jnp.array(batch["input_ids"])
batch_size = input_ids.shape[0]
label_drop_mask = jnp.zeros((batch_size,), dtype=jnp.bool_)
if config.label_drop_prob > 0:
rng, drop_rng = jax.random.split(rng)
label_drop_mask = jax.random.uniform(drop_rng, (batch_size,)) < config.label_drop_prob
result["label_drop_mask"] = label_drop_mask
return result
def pad_and_truncate(ids_list, target_len, pad_token_id):
"""Pad or truncate sequences to target_len, return stacked array and lengths."""
padded, lengths = [], []
for ids in ids_list:
orig_len = min(len(ids), target_len)
ids = ids[:target_len]
if orig_len < target_len:
ids = np.concatenate([ids, np.full(target_len - orig_len, pad_token_id, dtype=ids.dtype)])
padded.append(ids)
lengths.append(orig_len)
return np.stack(padded), np.array(lengths)
def get_dataloader(
dataset,
batch_size: int,
shuffle: bool = True,
num_workers: int = 0,
drop_last: bool = True,
max_seq_length: int = 512,
pad_token_id: int = 0,
max_input_seq_length: Optional[int] = None,
distributed: bool = True,
):
"""Create a DataLoader."""
def collate_fn(batch_list):
input_ids_list = [np.array(item["input_ids"]) for item in batch_list]
if "condition_input_ids" in batch_list[0]:
seq_list, cond_lens = [], []
for item in batch_list:
cond = np.array(item["condition_input_ids"])[:max_input_seq_length]
inp = np.array(item["input_ids"])
seq_list.append(np.concatenate([cond, inp]))
cond_lens.append(len(cond))
cond_lens = np.array(cond_lens)
else:
seq_list = input_ids_list
cond_lens = np.zeros(len(input_ids_list), dtype=np.int32)
ids, total_lens = pad_and_truncate(seq_list, max_seq_length, pad_token_id)
pos = np.arange(max_seq_length)[None, :]
is_cond = pos < cond_lens[:, None]
is_valid = pos < total_lens[:, None]
encoder_attn, attn, pred = build_self_attn_cond_masks(is_cond, is_valid, xp=np)
result = {
"input_ids": ids,
"encoder_attention_mask": encoder_attn,
"attention_mask": attn,
"cond_seq_mask": pred,
}
for key in ("index", "input", "target"):
if key in batch_list[0]:
result[key] = [item[key] for item in batch_list]
return result
common = dict(
batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn,
drop_last=drop_last, persistent_workers=num_workers > 0,
)
if distributed:
sampler = DistributedSampler(
dataset, num_replicas=jax.process_count(), rank=jax.process_index(),
shuffle=shuffle, drop_last=drop_last,
)
return DataLoader(dataset, sampler=sampler, **common)
return DataLoader(dataset, shuffle=shuffle, **common)
def load_jsonl_dataset(path, tokenizer, input_key="input", output_key="output"):
"""Load a JSONL eval set (one `{input, output}` example per line).
Triggered by `eval.py` whenever `config.eval_data_path` ends with `.jsonl`;
otherwise the standard `datasets.load_from_disk` Arrow path is used.
"""
examples = []
with open(path, "r", encoding="utf-8") as f:
for i, line in enumerate(f):
line = line.strip()
if not line:
continue
data = json.loads(line)
examples.append({
"index": i,
"input": data[input_key],
"target": data[output_key],
"condition_input_ids": tokenizer(data[input_key], add_special_tokens=False)["input_ids"],
"input_ids": tokenizer(data[output_key], add_special_tokens=False)["input_ids"],
})
return examples
# ============================================
# Dataset loading
# ============================================
def _looks_like_save_to_disk_arrow(ds) -> bool:
"""HF datasets uploaded via `save_to_disk` get loaded as a fake 1-row dataset
where the columns are internal metadata fields like `_data_files`, `_fingerprint`,
etc. Detect that here so we can fall back to `load_from_disk`."""
return (
len(ds) == 1
and any(c.startswith("_") for c in ds.column_names)
and not any(not c.startswith("_") for c in ds.column_names)
)
def load_dataset_split(path: str, dataset_cache_dir=None):
"""Load a dataset. Tries HuggingFace Hub first; falls back to local on-disk Arrow.
For HF repos uploaded via `dataset.save_to_disk()` instead of `push_to_hub()`,
`load_dataset` silently returns a 1-row dataset of internal metadata. We detect
that and re-download the repo, then load it via `load_from_disk`.
"""
ds = None
try:
ds = hf_load_dataset(path, cache_dir=dataset_cache_dir)
except Exception:
ds = load_from_disk(path)
# Normalize DatasetDict -> single split before detection.
if isinstance(ds, DatasetDict):
splits = list(ds.keys())
if len(splits) != 1:
raise ValueError(
f"Expected dataset at {path!r} to have a single split, got {splits}."
)
ds = ds[splits[0]]
# save_to_disk-style HF repo: fall back via snapshot_download + load_from_disk.
if _looks_like_save_to_disk_arrow(ds):
from huggingface_hub import snapshot_download
log_for_0(
f"Dataset at {path!r} looks like a save_to_disk-format HF repo; "
f"re-downloading via snapshot_download + load_from_disk."
)
local_dir = snapshot_download(
repo_id=path, repo_type="dataset", cache_dir=dataset_cache_dir,
)
ds = load_from_disk(local_dir)
if isinstance(ds, DatasetDict):
splits = list(ds.keys())
if len(splits) != 1:
raise ValueError(
f"Expected dataset at {path!r} to have a single split, got {splits}."
)
ds = ds[splits[0]]
ds.set_format(type="numpy", columns=ds.column_names)
return ds
def load_dataset(config, dataset_cache_dir=None):
"""Resolve config.data_path / config.eval_data_path into train/eval datasets."""
log_for_0(f"Loading dataset from {config.data_path}...")
train_dataset = load_dataset_split(config.data_path, dataset_cache_dir)
log_for_0(f"Train size: {len(train_dataset)}")
eval_dataset = None
if config.eval_data_path:
eval_dataset = load_dataset_split(config.eval_data_path, dataset_cache_dir)
log_for_0(f"Eval size: {len(eval_dataset)}")
else:
log_for_0("No eval dataset")
return train_dataset, eval_dataset
|