delta/dataset.py

"""Dataset and collator for Delta causal language modeling."""

from __future__ import annotations

import csv
import io
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 DEFAULT_SYSTEM_PROMPT, DeltaTokenizer

logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)

RAW_TEXT_SUFFIXES = {".txt", ".md", ".markdown"}
STRUCTURED_SUFFIXES = {".jsonl", ".json", ".csv"}
SUPPORTED_SUFFIXES = RAW_TEXT_SUFFIXES | STRUCTURED_SUFFIXES


def _read_file_text(path: Path) -> str:
    """Read text while tolerating common Windows/UTF-8 corpus encodings."""

    last_error: UnicodeDecodeError | None = None
    for encoding in ("utf-8-sig", "utf-8", "utf-16", "cp1252"):
        try:
            return path.read_text(encoding=encoding)
        except UnicodeDecodeError as exc:
            last_error = exc
    if last_error is not None:
        raise last_error
    return ""


def _mojibake_score(text: str) -> int:
    """Score common UTF-8-as-Windows-1252 artifacts."""

    markers = ("Ã", "Â", "â€", "�")
    return sum(text.count(marker) for marker in markers)


def _clean_text(text: str) -> str:
    """Normalize line endings and repair obvious mojibake when it improves text."""

    cleaned = text.replace("\r\n", "\n").replace("\r", "\n")
    if _mojibake_score(cleaned) == 0:
        return cleaned.strip()
    try:
        repaired = cleaned.encode("cp1252").decode("utf-8")
    except UnicodeError:
        return cleaned.strip()
    if _mojibake_score(repaired) < _mojibake_score(cleaned):
        return repaired.strip()
    return cleaned.strip()


def _iter_data_files(path: Path) -> list[Path]:
    """Return supported dataset files in stable order."""

    if path.is_file():
        return [path] if path.suffix.lower() in SUPPORTED_SUFFIXES else []
    files: list[Path] = []
    for file_path in sorted(path.rglob("*")):
        if not file_path.is_file():
            continue
        if file_path.name.startswith("."):
            continue
        if file_path.name.lower() == "readme.md":
            continue
        if file_path.suffix.lower() in SUPPORTED_SUFFIXES:
            files.append(file_path)
    return files


def _format_chat_messages(messages: list[Any], system: str | None = None) -> str:
    """Convert role/content messages into Delta chat-token training text."""

    parts: list[str] = []
    if system:
        parts.append(f"[SYS] {system.strip()} [SEP]")
    for message in messages:
        if not isinstance(message, dict):
            continue
        role = str(message.get("role", "")).lower().strip()
        content = str(message.get("content", "")).strip()
        if not content:
            continue
        if role == "system":
            if parts and parts[0].startswith("[SYS]"):
                parts[0] = f"[SYS] {content} [SEP]"
            else:
                parts.insert(0, f"[SYS] {content} [SEP]")
        elif role in {"user", "human", "prompt", "instruction"}:
            parts.append(f"[USR] {content} [SEP]")
        elif role in {"assistant", "model", "completion", "answer", "response"}:
            parts.append(f"[ASS] {content} [SEP]")
    if not parts or not parts[0].startswith("[SYS]"):
        parts.insert(0, f"[SYS] {DEFAULT_SYSTEM_PROMPT} [SEP]")
    return "\n".join(parts)


def _format_prompt_completion(record: dict[str, Any]) -> str | None:
    """Convert instruction/prompt datasets into Delta chat-token training text."""

    prompt = record.get("prompt") or record.get("question") or record.get("instruction")
    completion = (
        record.get("completion")
        or record.get("response")
        or record.get("answer")
        or record.get("output")
    )
    if prompt is None or completion is None:
        return None
    extra_input = str(record.get("input", "")).strip()
    user_text = str(prompt).strip()
    if extra_input:
        user_text = f"{user_text}\n\n{extra_input}"
    system = str(record.get("system") or DEFAULT_SYSTEM_PROMPT).strip()
    return "\n".join(
        [
            f"[SYS] {system} [SEP]",
            f"[USR] {user_text} [SEP]",
            f"[ASS] {str(completion).strip()} [SEP]",
        ]
    )


def _record_to_text(record: Any) -> str | None:
    """Convert a supported structured record into training text."""

    if isinstance(record, str):
        return record
    if not isinstance(record, dict):
        return None
    if "text" in record:
        return str(record["text"])
    if isinstance(record.get("messages"), list):
        return _format_chat_messages(record["messages"], system=record.get("system"))
    return _format_prompt_completion(record)


def _json_records(value: Any) -> list[Any]:
    """Extract records from common JSON dataset shapes."""

    if isinstance(value, list):
        return value
    if isinstance(value, dict):
        for key in ("data", "records", "examples", "samples"):
            if isinstance(value.get(key), list):
                return value[key]
        return [value]
    return []


def _read_jsonl(file_path: Path) -> list[str]:
    """Read JSONL records from a file."""

    texts: list[str] = []
    for line_number, line in enumerate(_read_file_text(file_path).splitlines(), start=1):
        line = line.strip()
        if not line:
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid JSONL in {file_path}:{line_number}: {exc}") from exc
        text = _record_to_text(record)
        if text:
            texts.append(text)
    return texts


def _read_json(file_path: Path) -> list[str]:
    """Read JSON records from object/list dataset files."""

    payload = json.loads(_read_file_text(file_path))
    return [text for record in _json_records(payload) if (text := _record_to_text(record))]


def _read_csv(file_path: Path) -> list[str]:
    """Read CSV datasets with text or prompt/completion-style columns."""

    texts: list[str] = []
    reader = csv.DictReader(io.StringIO(_read_file_text(file_path)))
    for row in reader:
        text = _record_to_text(row)
        if text:
            texts.append(text)
    return texts


def _read_texts(path: Path) -> list[str]:
    """Read supported dataset files into normalized training texts."""

    texts: list[str] = []
    files = _iter_data_files(path)
    for file_path in files:
        suffix = file_path.suffix.lower()
        if suffix in RAW_TEXT_SUFFIXES:
            texts.append(_read_file_text(file_path))
        elif suffix == ".jsonl":
            texts.extend(_read_jsonl(file_path))
        elif suffix == ".json":
            texts.extend(_read_json(file_path))
        elif suffix == ".csv":
            texts.extend(_read_csv(file_path))
    cleaned = [_clean_text(text) for text in texts]
    return [text for text in cleaned if text]


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 = 768,
        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:
            formats = ", ".join(sorted(SUPPORTED_SUFFIXES))
            raise ValueError(f"No supported dataset records ({formats}) 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}