| """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} |
|
|
