File size: 3,747 Bytes
a006332 | 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 | """Dataset and collator for Delta causal language modeling."""
from __future__ import annotations
import json
import logging
import os
from pathlib import Path
from typing import Any
import torch
from torch.utils.data import Dataset
from delta.tokenizer import DeltaTokenizer
logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)
def _read_texts(path: Path) -> list[str]:
"""Read .txt files or .jsonl files containing a text field."""
files = [path] if path.is_file() else sorted(path.rglob("*"))
texts: list[str] = []
for file_path in files:
if file_path.suffix.lower() == ".txt":
texts.append(file_path.read_text(encoding="utf-8"))
elif file_path.suffix.lower() == ".jsonl":
with file_path.open("r", encoding="utf-8") as handle:
for line in handle:
line = line.strip()
if not line:
continue
record = json.loads(line)
if "text" in record:
texts.append(str(record["text"]))
return texts
class DeltaDataset(Dataset[dict[str, torch.Tensor]]):
"""Sliding-window token dataset for language modeling."""
def __init__(
self,
data_path: str | Path,
tokenizer: DeltaTokenizer,
max_seq_len: int = 512,
stride: int = 256,
) -> None:
self.data_path = Path(data_path)
self.tokenizer = tokenizer
self.max_seq_len = max_seq_len
self.stride = stride
texts = _read_texts(self.data_path)
if not texts:
raise ValueError(f"No .txt or .jsonl text records found in {self.data_path}")
self.windows: list[list[int]] = []
for text in texts:
ids = tokenizer.encode(text, add_special_tokens=True)
for start in range(0, max(1, len(ids) - 1), stride):
window = ids[start : start + max_seq_len]
if len(window) >= 2:
self.windows.append(window)
if start + max_seq_len >= len(ids):
break
logger.info("Loaded %s training windows from %s", len(self.windows), self.data_path)
def __len__(self) -> int:
"""Return the number of windows."""
return len(self.windows)
def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
"""Return one token window."""
ids = torch.tensor(self.windows[index], dtype=torch.long)
return {"input_ids": ids, "labels": ids.clone()}
class DeltaDataCollator:
"""Dynamic padding collator for causal language modeling."""
def __init__(self, pad_token_id: int = 0) -> None:
self.pad_token_id = pad_token_id
def __call__(self, features: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
"""Pad input ids and labels to the longest sample in the batch."""
max_len = max(feature["input_ids"].size(0) for feature in features)
input_ids = torch.full((len(features), max_len), self.pad_token_id, dtype=torch.long)
labels = torch.full((len(features), max_len), -100, dtype=torch.long)
attention_mask = torch.zeros((len(features), max_len), dtype=torch.long)
for row, feature in enumerate(features):
ids = feature["input_ids"]
length = ids.size(0)
input_ids[row, :length] = ids
labels[row, :length] = feature["labels"]
pad_positions = ids == self.pad_token_id
labels[row, :length][pad_positions] = -100
attention_mask[row, :length] = 1
return {"input_ids": input_ids, "labels": labels, "attention_mask": attention_mask}
|