src/data/dataset.py

"""
Dataset classes for SLM training.

Handles loading, preprocessing, and tokenization of conversational data.
"""

import os
import json
import random
from typing import List, Dict, Optional, Iterator, Tuple
from pathlib import Path

import torch
from torch.utils.data import Dataset, IterableDataset

from .tokenizer import SLMTokenizer


class ConversationalDataset(Dataset):
    """Dataset for conversational/instruction-following data.

    Loads pre-tokenized data from disk for efficient training.
    Format: Each sample is a tokenized conversation with user/assistant turns.
    """

    def __init__(
        self,
        data_path: str,
        tokenizer: SLMTokenizer,
        max_length: int = 1024,
        split: str = "train",
    ):
        """Initialize the dataset.

        Args:
            data_path: Path to the processed data directory
            tokenizer: Tokenizer instance
            max_length: Maximum sequence length
            split: "train" or "val"
        """
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.split = split

        # Load data
        self.samples = self._load_data(data_path)
        print(f"Loaded {len(self.samples)} samples for {split} split")

    def _load_data(self, data_path: str) -> List[Dict]:
        """Load data from JSON or JSONL files."""
        samples = []

        # Check for split-specific JSONL file first (preferred for large datasets)
        split_jsonl = os.path.join(data_path, f"{self.split}.jsonl")
        if os.path.exists(split_jsonl):
            with open(split_jsonl, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        samples.append(json.loads(line))
            return samples

        # Check for split-specific JSON file
        split_file = os.path.join(data_path, f"{self.split}.json")
        if os.path.exists(split_file):
            with open(split_file, "r", encoding="utf-8") as f:
                # Try JSONL format first (one JSON per line)
                content = f.read()
                f.seek(0)
                try:
                    # Try loading as single JSON array
                    samples = json.loads(content)
                    if isinstance(samples, list):
                        return samples
                except json.JSONDecodeError:
                    pass

                # Load as JSONL (one JSON per line)
                for line in f:
                    line = line.strip()
                    if line:
                        samples.append(json.loads(line))
                return samples

        # Check for combined file with splits
        combined_file = os.path.join(data_path, "data.json")
        if os.path.exists(combined_file):
            with open(combined_file, "r") as f:
                all_data = json.load(f)
            if isinstance(all_data, dict) and self.split in all_data:
                return all_data[self.split]
            return all_data

        # Load all .json and .jsonl files in directory
        for ext in ["*.jsonl", "*.json"]:
            for file in sorted(Path(data_path).glob(ext)):
                with open(file, "r", encoding="utf-8") as f:
                    if file.suffix == ".jsonl":
                        for line in f:
                            line = line.strip()
                            if line:
                                samples.append(json.loads(line))
                    else:
                        data = json.load(f)
                        if isinstance(data, list):
                            samples.extend(data)
                        else:
                            samples.append(data)

        return samples

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Get a single sample.

        Returns:
            Dictionary with:
            - input_ids: Token IDs for the full sequence
            - attention_mask: 1 for real tokens, 0 for padding
            - labels: Same as input_ids but with -100 for padding (for loss)
        """
        sample = self.samples[idx]

        # Handle different data formats
        if "input_ids" in sample:
            # Pre-tokenized data
            input_ids = sample["input_ids"]
        elif "user" in sample and "assistant" in sample:
            # Raw conversation format
            input_ids = self.tokenizer.encode_conversation(
                user_message=sample["user"],
                assistant_message=sample["assistant"],
                max_length=self.max_length,
            )
        elif "text" in sample:
            # Raw text format
            input_ids = self.tokenizer.encode(
                sample["text"],
                add_special_tokens=True,
                max_length=self.max_length,
                truncation=True,
            )
        elif "question" in sample and "answer" in sample:
            # Q&A format
            input_ids = self.tokenizer.encode_conversation(
                user_message=sample["question"],
                assistant_message=sample["answer"],
                max_length=self.max_length,
            )
        else:
            raise ValueError(f"Unknown sample format: {list(sample.keys())}")

        # Pad or truncate
        if len(input_ids) > self.max_length:
            input_ids = input_ids[:self.max_length]
            # Ensure EOS at the end
            if input_ids[-1] != self.tokenizer.eos_token_id:
                input_ids[-1] = self.tokenizer.eos_token_id

        # Create attention mask (before padding)
        attention_mask = [1] * len(input_ids)

        # Pad if needed
        padding_length = self.max_length - len(input_ids)
        if padding_length > 0:
            input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_length
            attention_mask = attention_mask + [0] * padding_length

        # Labels for language modeling (shift happens in loss function)
        # Use -100 for padding tokens so they're ignored in loss
        labels = [
            id if mask == 1 else -100
            for id, mask in zip(input_ids, attention_mask)
        ]

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


class StreamingTextDataset(IterableDataset):
    """Streaming dataset for large text files.

    Memory-efficient dataset that streams data from disk.
    Useful for training on large text corpora.
    """

    def __init__(
        self,
        data_files: List[str],
        tokenizer: SLMTokenizer,
        max_length: int = 1024,
        shuffle: bool = True,
        seed: int = 42,
    ):
        """Initialize streaming dataset.

        Args:
            data_files: List of text file paths
            tokenizer: Tokenizer instance
            max_length: Maximum sequence length
            shuffle: Whether to shuffle files and lines
            seed: Random seed for shuffling
        """
        self.data_files = data_files
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.shuffle = shuffle
        self.seed = seed

        # Verify files exist
        for f in data_files:
            if not os.path.exists(f):
                raise FileNotFoundError(f"Data file not found: {f}")

    def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
        """Iterate over all samples in all files."""
        worker_info = torch.utils.data.get_worker_info()

        # Handle multi-worker data loading
        if worker_info is None:
            files_to_process = self.data_files
        else:
            # Split files among workers
            per_worker = len(self.data_files) // worker_info.num_workers
            worker_id = worker_info.id
            start = worker_id * per_worker
            end = start + per_worker if worker_id < worker_info.num_workers - 1 else len(self.data_files)
            files_to_process = self.data_files[start:end]

        # Shuffle files if needed
        if self.shuffle:
            rng = random.Random(self.seed)
            files_to_process = list(files_to_process)
            rng.shuffle(files_to_process)

        # Buffer for accumulating text
        buffer = []
        buffer_tokens = 0

        for file_path in files_to_process:
            with open(file_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    # Try to parse as JSON (for conversational data)
                    try:
                        data = json.loads(line)
                        if "user" in data and "assistant" in data:
                            tokens = self.tokenizer.encode_conversation(
                                data["user"], data["assistant"]
                            )
                        elif "text" in data:
                            tokens = self.tokenizer.encode(
                                data["text"], add_special_tokens=True
                            )
                        else:
                            tokens = self.tokenizer.encode(
                                line, add_special_tokens=True
                            )
                    except json.JSONDecodeError:
                        # Plain text line
                        tokens = self.tokenizer.encode(
                            line, add_special_tokens=True
                        )

                    buffer.extend(tokens)

                    # Yield chunks of max_length
                    while len(buffer) >= self.max_length:
                        chunk = buffer[:self.max_length]
                        buffer = buffer[self.max_length:]

                        yield self._create_sample(chunk)

        # Handle remaining buffer (pad to max_length)
        if len(buffer) > 0:
            yield self._create_sample(buffer)

    def _create_sample(self, tokens: List[int]) -> Dict[str, torch.Tensor]:
        """Create a training sample from tokens."""
        input_ids = tokens[:self.max_length]

        # Pad if needed
        attention_mask = [1] * len(input_ids)
        padding_length = self.max_length - len(input_ids)
        if padding_length > 0:
            input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_length
            attention_mask = attention_mask + [0] * padding_length

        labels = [
            id if mask == 1 else -100
            for id, mask in zip(input_ids, attention_mask)
        ]

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


class PackedDataset(Dataset):
    """Dataset that packs multiple short sequences into one.

    Efficient for training when samples are shorter than max_length.
    Concatenates samples with separator tokens to fill sequences.
    """

    def __init__(
        self,
        samples: List[Dict],
        tokenizer: SLMTokenizer,
        max_length: int = 1024,
    ):
        """Initialize packed dataset.

        Args:
            samples: List of samples with "user" and "assistant" keys
            tokenizer: Tokenizer instance
            max_length: Maximum sequence length
        """
        self.tokenizer = tokenizer
        self.max_length = max_length

        # Pack sequences
        self.packed_samples = self._pack_sequences(samples)
        print(f"Packed {len(samples)} samples into {len(self.packed_samples)} sequences")

    def _pack_sequences(self, samples: List[Dict]) -> List[List[int]]:
        """Pack short sequences together."""
        packed = []
        current_sequence = []

        for sample in samples:
            # Tokenize
            if "user" in sample and "assistant" in sample:
                tokens = self.tokenizer.encode_conversation(
                    sample["user"], sample["assistant"]
                )
            elif "text" in sample:
                tokens = self.tokenizer.encode(sample["text"], add_special_tokens=True)
            else:
                continue

            # Check if we can add to current sequence
            if len(current_sequence) + len(tokens) <= self.max_length:
                current_sequence.extend(tokens)
            else:
                # Save current and start new
                if current_sequence:
                    packed.append(current_sequence)
                current_sequence = tokens[:self.max_length]

        # Don't forget the last sequence
        if current_sequence:
            packed.append(current_sequence)

        return packed

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

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Get a packed sample."""
        tokens = self.packed_samples[idx]

        # Pad if needed
        attention_mask = [1] * len(tokens)
        padding_length = self.max_length - len(tokens)
        if padding_length > 0:
            tokens = tokens + [self.tokenizer.pad_token_id] * padding_length
            attention_mask = attention_mask + [0] * padding_length

        labels = [
            id if mask == 1 else -100
            for id, mask in zip(tokens, attention_mask)
        ]

        return {
            "input_ids": torch.tensor(tokens, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


def create_train_val_split(
    samples: List[Dict],
    val_ratio: float = 0.01,
    seed: int = 42,
) -> Tuple[List[Dict], List[Dict]]:
    """Split samples into train and validation sets.

    Args:
        samples: List of all samples
        val_ratio: Ratio for validation set
        seed: Random seed

    Returns:
        Tuple of (train_samples, val_samples)
    """
    random.seed(seed)
    shuffled = list(samples)
    random.shuffle(shuffled)

    val_size = int(len(shuffled) * val_ratio)
    val_samples = shuffled[:val_size]
    train_samples = shuffled[val_size:]

    return train_samples, val_samples


def load_jsonl(file_path: str) -> List[Dict]:
    """Load data from a JSONL file."""
    samples = []
    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                samples.append(json.loads(line))
    return samples


def save_jsonl(samples: List[Dict], file_path: str):
    """Save data to a JSONL file."""
    with open(file_path, "w", encoding="utf-8") as f:
        for sample in samples:
            f.write(json.dumps(sample) + "\n")