src/training_data.py

from __future__ import annotations

import json
import os

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset

from configs import cfg

PAD_MULTIPLE = 128

def torch_load_cpu(path: str) -> dict:
    try:
        return torch.load(path, map_location="cpu", weights_only=True)
    except TypeError:
        return torch.load(path, map_location="cpu")

def extract_shard_id_range(shard_payload: dict, shard_path: str) -> tuple[int, int]:
    try:
        ids_payload = shard_payload["ids"]
    except KeyError as exc:
        raise KeyError(
            f"Teacher shard {shard_path} is missing 'ids'. Regenerate the teacher-logit shards."
        ) from exc

    if torch.is_tensor(ids_payload):
        if ids_payload.numel() == 0:
            raise ValueError(
                f"Teacher shard {shard_path} has an empty ids tensor. Regenerate the teacher-logit shards."
            )
        return int(ids_payload.min().item()), int(ids_payload.max().item())

    if not isinstance(ids_payload, list) or not ids_payload:
        raise ValueError(
            f"Teacher shard {shard_path} has an incompatible ids payload. "
            "Regenerate the teacher-logit shards."
        )

    min_id: int | None = None
    max_id: int | None = None
    for sample_idx, ids_tensor in enumerate(ids_payload):
        if not torch.is_tensor(ids_tensor):
            raise ValueError(
                f"Teacher shard {shard_path} sample #{sample_idx} has a non-tensor ids payload. "
                "Regenerate the teacher-logit shards."
            )
        if ids_tensor.numel() == 0:
            continue
        sample_min = int(ids_tensor.min().item())
        sample_max = int(ids_tensor.max().item())
        min_id = sample_min if min_id is None else min(min_id, sample_min)
        max_id = sample_max if max_id is None else max(max_id, sample_max)

    if min_id is None or max_id is None:
        raise ValueError(
            f"Teacher shard {shard_path} only contains empty ids tensors. "
            "Regenerate the teacher-logit shards."
        )
    return min_id, max_id

class DistillationDataset(Dataset):
    def __init__(self, data_path: str, logits_dir: str, max_seq_len: int, num_samples: int = -1, phase: str = "kd"):
        self.phase = phase
        self.data_path = data_path
        self.logits_dir = logits_dir
        self.max_seq_len = max_seq_len
        self.samples_per_shard = self._resolve_samples_per_shard()
        self.sample_offsets: list[int] = []
        self.sample_lengths: list[int] = []
        self.sample_target_counts: list[int] = []
        self._data_handle = None
        self._cached_shard_idx: int | None = None
        self._cached_shard_path: str | None = None
        self._cached_shard_payload: dict | None = None

        with open(data_path, "r", encoding="utf-8") as f:
            while True:
                if 0 < num_samples <= len(self.sample_offsets):
                    break
                offset = f.tell()
                line = f.readline()
                if not line:
                    break
                i = len(self.sample_offsets)
                raw_sample = json.loads(line)
                input_ids_list, loss_mask_list = self._coerce_tokenized_row(raw_sample, i)
                self.sample_offsets.append(offset)
                self.sample_lengths.append(len(input_ids_list))
                self.sample_target_counts.append(sum(loss_mask_list))

    def __len__(self) -> int:
        return len(self.sample_offsets)

    def __getstate__(self) -> dict:
        state = self.__dict__.copy()
        state["_data_handle"] = None
        state["_cached_shard_idx"] = None
        state["_cached_shard_path"] = None
        state["_cached_shard_payload"] = None
        return state

    def __del__(self) -> None:
        data_handle = getattr(self, "_data_handle", None)
        if data_handle is not None:
            try:
                data_handle.close()
            except Exception:
                pass

    def _resolve_samples_per_shard(self) -> int:
        prov_path = os.path.join(self.logits_dir, "_provenance.json")
        if not os.path.exists(prov_path):
            return 1

        try:
            with open(prov_path, "r", encoding="utf-8") as f:
                prov = json.load(f)
        except (OSError, json.JSONDecodeError):
            return 1

        shard_schema = prov.get("shard_schema", {})
        if shard_schema.get("layout") != "chunked_sample_lists":
            return 1

        raw_value = prov.get("samples_per_shard", 1)
        try:
            value = int(raw_value)
        except (TypeError, ValueError):
            return 1
        return max(value, 1)

    def _coerce_tokenized_row(self, raw_sample: dict, idx: int) -> tuple[list[int], list[int]]:
        try:
            input_ids = raw_sample["input_ids"][: self.max_seq_len]
        except KeyError as exc:
            raise KeyError(
                f"Tokenized sample #{idx} is missing 'input_ids'. "
                "Re-run download.py to regenerate the tokenized dataset."
            ) from exc
        try:
            loss_mask = raw_sample["loss_mask"][: len(input_ids)]
        except KeyError as exc:
            raise KeyError(
                "Tokenized sample is missing 'loss_mask'. Re-run download.py to regenerate "
                "assistant-only training targets before distilling."
            ) from exc

        if not isinstance(input_ids, list) or len(input_ids) == 0:
            raise ValueError(
                f"Tokenized sample #{idx} has incompatible input_ids payload. "
                "Re-run download.py to regenerate."
            )
        if not isinstance(loss_mask, list) or len(loss_mask) != len(input_ids):
            raise ValueError(
                f"Tokenized sample #{idx} has incompatible loss_mask length {len(loss_mask)}. "
                "Re-run download.py to regenerate assistant-only targets."
            )

        normalized_mask = [int(value) for value in loss_mask]
        if any(value not in (0, 1) for value in normalized_mask):
            raise ValueError(
                f"Tokenized sample #{idx} has non-binary loss_mask values. "
                "Re-run download.py to regenerate assistant-only targets."
            )
        if sum(normalized_mask) == 0:
            raise ValueError(
                f"Tokenized sample #{idx} has no assistant target tokens. "
                "Re-run download.py to filter invalid conversations."
            )

        return [int(token_id) for token_id in input_ids], normalized_mask

    def _data_file(self):
        if self._data_handle is None:
            self._data_handle = open(self.data_path, "r", encoding="utf-8")
        return self._data_handle

    def _load_raw_sample(self, idx: int) -> dict:
        data_file = self._data_file()
        data_file.seek(self.sample_offsets[idx])
        line = data_file.readline()
        if not line:
            raise IndexError(f"Tokenized sample #{idx} could not be read from {self.data_path}.")
        return json.loads(line)

    def _load_shard_payload(self, shard_idx: int) -> tuple[str, dict]:
        if self._cached_shard_idx == shard_idx and self._cached_shard_payload is not None and self._cached_shard_path is not None:
            return self._cached_shard_path, self._cached_shard_payload

        shard_path = os.path.join(self.logits_dir, f"shard_{shard_idx:06d}.pt")
        if not os.path.exists(shard_path):
            raise FileNotFoundError(
                f"Missing teacher logit shard: {shard_path}. "
                "Regenerate the teacher-logit shards."
            )

        payload = torch_load_cpu(shard_path)
        self._cached_shard_idx = shard_idx
        self._cached_shard_path = shard_path
        self._cached_shard_payload = payload
        return shard_path, payload

    def _load_teacher_tensors(self, idx: int, seq_len: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if self.samples_per_shard <= 1:
            shard_path, shard = self._load_shard_payload(idx)
            try:
                teacher_logprobs = shard["logprobs"][:seq_len]
                teacher_ids = shard["ids"][:seq_len]
                teacher_other_logprob = shard["other_logprob"][:seq_len]
            except KeyError as exc:
                missing = exc.args[0]
                raise KeyError(
                    f"Shard {shard_path} is missing {missing!r}. "
                    "Regenerate the current teacher-logit shards."
                ) from exc
            return teacher_logprobs, teacher_ids, teacher_other_logprob

        shard_idx = idx // self.samples_per_shard
        sample_offset = idx % self.samples_per_shard
        shard_path, shard = self._load_shard_payload(shard_idx)
        try:
            count = int(shard["count"])
            start_idx = int(shard["start_idx"])
            logprobs_list = shard["logprobs"]
            ids_list = shard["ids"]
            other_list = shard["other_logprob"]
        except KeyError as exc:
            missing = exc.args[0]
            raise KeyError(
                f"Grouped shard {shard_path} is missing {missing!r}. "
                "Regenerate the current teacher-logit shards."
            ) from exc

        expected_start_idx = shard_idx * self.samples_per_shard
        if start_idx != expected_start_idx:
            raise ValueError(
                f"Grouped shard {shard_path} starts at sample {start_idx}, "
                f"expected {expected_start_idx}. Regenerate the teacher-logit shards."
            )
        if not (len(logprobs_list) == len(ids_list) == len(other_list) == count):
            raise ValueError(
                f"Grouped shard {shard_path} has inconsistent sample counts. "
                "Regenerate the current teacher-logit shards."
            )
        if sample_offset >= count:
            raise FileNotFoundError(
                f"Grouped shard {shard_path} does not contain sample #{idx} "
                f"(start_idx={start_idx}, count={count}). Regenerate the teacher-logit shards."
            )
        try:
            teacher_logprobs = logprobs_list[sample_offset][:seq_len]
            teacher_ids = ids_list[sample_offset][:seq_len]
            teacher_other_logprob = other_list[sample_offset][:seq_len]
        except (IndexError, TypeError) as exc:
            raise ValueError(
                f"Grouped shard {shard_path} has an incompatible payload layout. "
                "Regenerate the current teacher-logit shards."
            ) from exc
        return teacher_logprobs, teacher_ids, teacher_other_logprob

    def __getitem__(self, idx: int) -> dict:
        raw_sample = self._load_raw_sample(idx)
        input_ids_list, loss_mask_list = self._coerce_tokenized_row(raw_sample, idx)
        input_ids = torch.tensor(input_ids_list, dtype=torch.long)
        loss_mask = torch.tensor(loss_mask_list, dtype=torch.long)
        seq_len = int(input_ids.size(0))

        if self.phase in ("sft", "online_kd"):
            return {"input_ids": input_ids, "loss_mask": loss_mask}

        teacher_logprobs, teacher_ids, teacher_other_logprob = self._load_teacher_tensors(idx, seq_len)
        if teacher_logprobs.shape[0] != seq_len:
            raise ValueError(
                f"Teacher shard for sample #{idx} has length {teacher_logprobs.shape[0]}, "
                f"but the tokenized row has length {seq_len}. Regenerate the teacher-logit shards; "
                "teacher shards must be in original JSONL row order."
            )
        if teacher_logprobs.ndim != 2 or teacher_ids.shape != teacher_logprobs.shape:
            raise ValueError(
                f"Teacher shard for sample #{idx} has incompatible top-k tensor shapes: "
                f"logprobs={tuple(teacher_logprobs.shape)}, ids={tuple(teacher_ids.shape)}. "
                "Regenerate the current teacher-logit shards."
            )
        if teacher_other_logprob.ndim != 1 or teacher_other_logprob.shape[0] != teacher_logprobs.shape[0]:
            raise ValueError(
                f"Teacher shard for sample #{idx} has incompatible other-bucket shape: "
                f"other_logprob={tuple(teacher_other_logprob.shape)}, "
                f"expected ({teacher_logprobs.shape[0]},). "
                "Regenerate the current teacher-logit shards."
            )
        if teacher_logprobs.shape[1] != cfg.training.top_k:
            raise ValueError(
                f"Teacher shard for sample #{idx} stores top_k={teacher_logprobs.shape[1]}, "
                f"expected {cfg.training.top_k}. "
                "Regenerate compatible teacher-logit shards."
            )

        return {
            "input_ids": input_ids,
            "loss_mask": loss_mask,
            "teacher_logprobs": teacher_logprobs,
            "teacher_ids": teacher_ids.long(),
            "teacher_other_logprob": teacher_other_logprob,
        }

def collate_fn(batch: list[dict], pad_token_id: int) -> dict:
    raw_max = max(item["input_ids"].size(0) for item in batch)
    max_len = ((raw_max + PAD_MULTIPLE - 1) // PAD_MULTIPLE) * PAD_MULTIPLE

    input_ids_list, mask_list, loss_mask_list, labels_list = [], [], [], []
    teacher_logprobs_list, teacher_ids_list, teacher_other_logprob_list = [], [], []

    for item in batch:
        seq_len = item["input_ids"].size(0)
        pad_len = max_len - seq_len
        padded_loss_mask = F.pad(item["loss_mask"], (0, pad_len), value=0)
        padded_labels = F.pad(item["input_ids"].clone(), (0, pad_len), value=pad_token_id)
        padded_labels = padded_labels.masked_fill(padded_loss_mask == 0, -100)

        input_ids_list.append(F.pad(item["input_ids"], (0, pad_len), value=pad_token_id))
        mask_list.append(
            torch.cat(
                [
                    torch.ones(seq_len, dtype=torch.long),
                    torch.zeros(pad_len, dtype=torch.long),
                ]
            )
        )
        loss_mask_list.append(padded_loss_mask)
        labels_list.append(padded_labels)

        if "teacher_logprobs" in item:
            teacher_seq_len = item["teacher_logprobs"].size(0)
            teacher_pad_len = max_len - teacher_seq_len
            teacher_logprobs_list.append(
                F.pad(item["teacher_logprobs"], (0, 0, 0, teacher_pad_len), value=float("-inf"))
            )
            teacher_ids_list.append(F.pad(item["teacher_ids"], (0, 0, 0, teacher_pad_len), value=0))
            teacher_other_logprob_list.append(
                F.pad(item["teacher_other_logprob"], (0, teacher_pad_len), value=float("-inf"))
            )

    result = {
        "input_ids": torch.stack(input_ids_list),
        "attention_mask": torch.stack(mask_list),
        "loss_mask": torch.stack(loss_mask_list).long(),
        "labels": torch.stack(labels_list),
    }
    if teacher_logprobs_list:
        result["teacher_logprobs"] = torch.stack(teacher_logprobs_list)
        result["teacher_ids"] = torch.stack(teacher_ids_list)
        result["teacher_other_logprob"] = torch.stack(teacher_other_logprob_list)
    return result

def resolve_dataloader_runtime() -> dict[str, int | bool]:
    cpu_count = max(1, os.cpu_count() or 1)
    configured_workers = int(getattr(cfg.training, "dataloader_workers", 4))
    num_workers = max(0, min(configured_workers, cpu_count))
    runtime: dict[str, int | bool] = {
        "num_workers": num_workers,
        "pin_memory": torch.cuda.is_available(),
    }
    if num_workers > 0:
        runtime["persistent_workers"] = True
        runtime["prefetch_factor"] = max(1, int(getattr(cfg.training, "prefetch_factor", 2)))
    return runtime

def move_batch_to_device(batch: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]:
    non_blocking = device.type == "cuda"
    return {
        name: tensor.to(device, non_blocking=non_blocking)
        for name, tensor in batch.items()
    }