Upload folder using huggingface_hub

187a491 verified 4 months ago

12.7 kB

"""
BoccaccioAI - PyTorch Datasets for Pre-Training and Instruction Fine-Tuning
Provides Dataset classes for loading pre-tokenized binary data (pre-training)
and JSONL instruction data (fine-tuning).
De Lauretis tech
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from tokenizers import Tokenizer
from torch.utils.data import Dataset


class PreTokenizedDataset(Dataset):
    """Dataset for loading pre-tokenized binary data during pre-training.
    Reads a flat uint16 numpy memmap file produced by ``tokenize_corpus.py``
    and serves fixed-length token sequences.
    """

    def __init__(self, data_path: str, sequence_length: int = 2048) -> None:
        self.data_path = Path(data_path)
        self.sequence_length = sequence_length

        if not self.data_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.data_path}")

        self.data = np.memmap(self.data_path, dtype=np.uint16, mode="r")
        self.num_sequences = len(self.data) // self.sequence_length

        if self.num_sequences == 0:
            raise ValueError(
                f"Data file {self.data_path} contains {len(self.data)} tokens, "
                f"which is fewer than one full sequence of length {self.sequence_length}."
            )

    def __len__(self) -> int:
        return self.num_sequences

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        if idx < 0 or idx >= self.num_sequences:
            raise IndexError(
                f"Index {idx} out of range for dataset with {self.num_sequences} sequences"
            )

        start = idx * self.sequence_length
        end = start + self.sequence_length
        tokens = np.array(self.data[start:end], dtype=np.int64)
        tokens = torch.from_numpy(tokens)

        return {"input_ids": tokens, "labels": tokens}


class InstructionDataset(Dataset):
    """Dataset for instruction fine-tuning with multi-turn conversations.
    Reads JSONL files with the standard format::
        {"conversations": [
            {"role": "system", "content": "..."},     # optional
            {"role": "user", "content": "..."},
            {"role": "assistant", "content": "..."},
            ...
        ]}
    Each conversation is rendered with explicit role markers::
        ### System:
        {system_content}
        ### Utente:
        {user_content}
        ### Assistente:
        {assistant_content}</s>
        ### Utente:
        {user_content_2}
        ### Assistente:
        {assistant_content_2}</s>
    Loss is computed only on assistant tokens (plus the trailing </s>).
    All user, system, and marker tokens are masked with -100.
    Examples longer than ``sequence_length`` after tokenization are skipped
    during ``__init__`` to avoid truncating responses (which would prevent
    the model from learning the EOS token).
    """

    SYSTEM_MARKER = "### System:"
    USER_MARKER = "### Utente:"
    ASSISTANT_MARKER = "### Assistente:"

    def __init__(
        self,
        data_path: str,
        tokenizer_path: str,
        sequence_length: int = 2048,
    ) -> None:
        self.data_path = Path(data_path)
        self.sequence_length = sequence_length

        if not self.data_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.data_path}")

        # Load tokenizer.
        tokenizer_file = Path(tokenizer_path)
        if not tokenizer_file.exists():
            raise FileNotFoundError(f"Tokenizer not found: {tokenizer_file}")
        self.tokenizer = Tokenizer.from_file(str(tokenizer_file))

        # Resolve special token IDs.
        self.pad_token_id = self.tokenizer.token_to_id("<pad>")
        if self.pad_token_id is None:
            self.pad_token_id = 0
        self.eos_token_id = self.tokenizer.token_to_id("</s>")
        if self.eos_token_id is None:
            raise ValueError("EOS token '</s>' not found in tokenizer vocabulary")

        # Load raw examples (JSONL). Tokenization happens on-the-fly in __getitem__
        # to avoid storing 268K pre-tokenized tensors in memory (~13 GB).
        self.raw_examples: list[dict] = []
        with open(self.data_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                self.raw_examples.append(json.loads(line))

        if not self.raw_examples:
            raise ValueError(f"No examples found in {self.data_path}")

        # Pre-compute token lengths to filter out examples that exceed sequence_length.
        # This is fast (just encoding, no tensor creation) and avoids wasting time
        # on examples that would be skipped anyway.
        valid_indices: list[int] = []
        skipped = 0
        for i, ex in enumerate(self.raw_examples):
            token_len = self._estimate_token_length(ex)
            if token_len <= self.sequence_length and token_len >= 10:
                valid_indices.append(i)
            else:
                skipped += 1

        if not valid_indices:
            raise ValueError(
                f"All examples were skipped (too long for sequence_length={sequence_length}). "
                f"Consider increasing sequence_length."
            )

        if skipped > 0:
            print(
                f"  InstructionDataset: skipped {skipped} examples "
                f"(exceeded sequence_length={sequence_length})"
            )

        self.valid_indices = valid_indices

    def _render_turn(self, role: str, content: str) -> str:
        """Render a single conversation turn with its role marker."""
        if role == "system":
            return f"{self.SYSTEM_MARKER}\n{content}\n"
        elif role == "user":
            return f"{self.USER_MARKER}\n{content}\n"
        elif role == "assistant":
            return f"{self.ASSISTANT_MARKER}\n{content}</s>\n"
        else:
            raise ValueError(f"Unknown role: {role}")

    def _estimate_token_length(self, ex: dict) -> int:
        """Exact token count for filtering.
        Mirrors the per-turn tokenization of _process_example exactly: BPE is
        not additive across turn boundaries, so encoding the concatenated text
        in one pass can undercount versus per-turn encoding.
        """
        conversations = ex.get("conversations", [])
        if not conversations or len(conversations) < 2:
            return 0
        if conversations[-1]["role"] != "assistant":
            return 0
        total = 0
        for turn in conversations:
            role = turn.get("role", "")
            content = (turn.get("content") or "").strip()
            if not content:
                continue
            if role in ("system", "user"):
                text = self._render_turn(role, content)
                total += len(self.tokenizer.encode(text).ids)
            elif role == "assistant":
                marker_text = f"{self.ASSISTANT_MARKER}\n"
                total += len(self.tokenizer.encode(marker_text).ids)
                content_text = f"{content}</s>\n"
                total += len(self.tokenizer.encode(content_text).ids)
        return total

    def _process_example(self, ex: dict) -> dict[str, torch.Tensor] | None:
        """Tokenize a conversation example and build labels with masking.
        Returns None if the example is too long for sequence_length.
        """
        conversations = ex.get("conversations", [])
        if not conversations or len(conversations) < 2:
            return None

        # Validate: must end with assistant turn.
        if conversations[-1]["role"] != "assistant":
            return None

        input_ids: list[int] = []
        labels: list[int] = []  # -100 for masked, token_id for assistant content + </s>

        for turn in conversations:
            role = turn.get("role", "")
            content = (turn.get("content") or "").strip()
            if not content:
                continue

            if role == "system":
                # Render and tokenize the marker + content. All masked.
                text = self._render_turn(role, content)
                ids = self.tokenizer.encode(text).ids
                input_ids.extend(ids)
                labels.extend([-100] * len(ids))

            elif role == "user":
                # Render and tokenize the marker + content. All masked.
                text = self._render_turn(role, content)
                ids = self.tokenizer.encode(text).ids
                input_ids.extend(ids)
                labels.extend([-100] * len(ids))

            elif role == "assistant":
                # Render marker (masked) + content + </s> (not masked).
                marker_text = f"{self.ASSISTANT_MARKER}\n"
                marker_ids = self.tokenizer.encode(marker_text).ids

                content_text = f"{content}</s>\n"
                content_ids = self.tokenizer.encode(content_text).ids

                input_ids.extend(marker_ids)
                labels.extend([-100] * len(marker_ids))

                input_ids.extend(content_ids)
                labels.extend(content_ids)

            else:
                # Unknown role, skip.
                continue

        # Check length: skip if too long (don't truncate - would lose </s>).
        if len(input_ids) > self.sequence_length:
            return None
        if len(input_ids) < 10:
            return None

        actual_length = len(input_ids)
        padding_length = self.sequence_length - actual_length

        # Pad input_ids with pad_token_id.
        input_ids_padded = input_ids + [self.pad_token_id] * padding_length

        # Pad labels with -100 (ignored in loss).
        labels_padded = labels + [-100] * padding_length

        # Attention mask: 1 for real tokens, 0 for padding.
        attention_mask = [1] * actual_length + [0] * padding_length

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

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

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        if idx < 0 or idx >= len(self.valid_indices):
            raise IndexError(
                f"Index {idx} out of range for dataset with {len(self.valid_indices)} examples"
            )
        raw_idx = self.valid_indices[idx]
        ex = self.raw_examples[raw_idx]
        processed = self._process_example(ex)
        # processed should not be None since we pre-filtered, but guard anyway.
        if processed is None:
            # Fallback: return a minimal valid example (shouldn't happen).
            raise RuntimeError(
                f"Example at index {idx} (raw {raw_idx}) failed to process "
                f"despite passing pre-filter. This is a bug."
            )
        return processed


if __name__ == "__main__":
    import sys

    print("BoccaccioAI Dataset classes")
    print(f"  PreTokenizedDataset: for pre-training on binary memmap files")
    print(f"  InstructionDataset:  for instruction fine-tuning on JSONL data")

    if len(sys.argv) > 1 and sys.argv[1] == "--test-pretokenized":
        if len(sys.argv) < 3:
            print("Usage: python dataset.py --test-pretokenized <data_path>")
            sys.exit(1)
        ds = PreTokenizedDataset(sys.argv[2])
        print(f"  Loaded {len(ds)} sequences from {sys.argv[2]}")
        sample = ds[0]
        print(f"  Sample input_ids shape: {sample['input_ids'].shape}")
        print(f"  Sample input_ids dtype: {sample['input_ids'].dtype}")
        print(f"  First 20 token IDs: {sample['input_ids'][:20].tolist()}")

    if len(sys.argv) > 1 and sys.argv[1] == "--test-instruction":
        if len(sys.argv) < 4:
            print("Usage: python dataset.py --test-instruction <data_path> <tokenizer_path>")
            sys.exit(1)
        ds = InstructionDataset(sys.argv[2], sys.argv[3])
        print(f"  Loaded {len(ds)} examples from {sys.argv[2]}")
        sample = ds[0]
        print(f"  Sample input_ids shape: {sample['input_ids'].shape}")
        print(f"  Sample labels shape:    {sample['labels'].shape}")
        print(f"  Sample attn_mask shape: {sample['attention_mask'].shape}")
        num_masked = (sample["labels"] == -100).sum().item()
        num_active = sample["labels"].numel() - num_masked
        print(f"  Masked label positions: {num_masked}")
        print(f"  Active label positions: {num_active} (assistant tokens + </s>)")

	"""
	BoccaccioAI - PyTorch Datasets for Pre-Training and Instruction Fine-Tuning

	Provides Dataset classes for loading pre-tokenized binary data (pre-training)
	and JSONL instruction data (fine-tuning).

	De Lauretis tech
	"""

	from __future__ import annotations

	import json
	from pathlib import Path

	import numpy as np
	import torch
	from tokenizers import Tokenizer
	from torch.utils.data import Dataset


	class PreTokenizedDataset(Dataset):
	"""Dataset for loading pre-tokenized binary data during pre-training.

	Reads a flat uint16 numpy memmap file produced by ``tokenize_corpus.py``
	and serves fixed-length token sequences.
	"""

	def __init__(self, data_path: str, sequence_length: int = 2048) -> None:
	self.data_path = Path(data_path)
	self.sequence_length = sequence_length

	if not self.data_path.exists():
	raise FileNotFoundError(f"Data file not found: {self.data_path}")

	self.data = np.memmap(self.data_path, dtype=np.uint16, mode="r")
	self.num_sequences = len(self.data) // self.sequence_length

	if self.num_sequences == 0:
	raise ValueError(
	f"Data file {self.data_path} contains {len(self.data)} tokens, "
	f"which is fewer than one full sequence of length {self.sequence_length}."
	)

	def __len__(self) -> int:
	return self.num_sequences

	def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
	if idx < 0 or idx >= self.num_sequences:
	raise IndexError(
	f"Index {idx} out of range for dataset with {self.num_sequences} sequences"
	)

	start = idx * self.sequence_length
	end = start + self.sequence_length
	tokens = np.array(self.data[start:end], dtype=np.int64)
	tokens = torch.from_numpy(tokens)

	return {"input_ids": tokens, "labels": tokens}


	class InstructionDataset(Dataset):
	"""Dataset for instruction fine-tuning with multi-turn conversations.

	Reads JSONL files with the standard format::

	{"conversations": [
	{"role": "system", "content": "..."}, # optional
	{"role": "user", "content": "..."},
	{"role": "assistant", "content": "..."},
	...
	]}

	Each conversation is rendered with explicit role markers::

	### System:
	{system_content}
	### Utente:
	{user_content}
	### Assistente:
	{assistant_content}</s>
	### Utente:
	{user_content_2}
	### Assistente:
	{assistant_content_2}</s>

	Loss is computed only on assistant tokens (plus the trailing </s>).
	All user, system, and marker tokens are masked with -100.

	Examples longer than ``sequence_length`` after tokenization are skipped
	during ``__init__`` to avoid truncating responses (which would prevent
	the model from learning the EOS token).
	"""

	SYSTEM_MARKER = "### System:"
	USER_MARKER = "### Utente:"
	ASSISTANT_MARKER = "### Assistente:"

	def __init__(
	self,
	data_path: str,
	tokenizer_path: str,
	sequence_length: int = 2048,
	) -> None:
	self.data_path = Path(data_path)
	self.sequence_length = sequence_length

	if not self.data_path.exists():
	raise FileNotFoundError(f"Data file not found: {self.data_path}")

	# Load tokenizer.
	tokenizer_file = Path(tokenizer_path)
	if not tokenizer_file.exists():
	raise FileNotFoundError(f"Tokenizer not found: {tokenizer_file}")
	self.tokenizer = Tokenizer.from_file(str(tokenizer_file))

	# Resolve special token IDs.
	self.pad_token_id = self.tokenizer.token_to_id("<pad>")
	if self.pad_token_id is None:
	self.pad_token_id = 0
	self.eos_token_id = self.tokenizer.token_to_id("</s>")
	if self.eos_token_id is None:
	raise ValueError("EOS token '</s>' not found in tokenizer vocabulary")

	# Load raw examples (JSONL). Tokenization happens on-the-fly in __getitem__
	# to avoid storing 268K pre-tokenized tensors in memory (~13 GB).
	self.raw_examples: list[dict] = []
	with open(self.data_path, "r", encoding="utf-8") as f:
	for line in f:
	line = line.strip()
	if not line:
	continue
	self.raw_examples.append(json.loads(line))

	if not self.raw_examples:
	raise ValueError(f"No examples found in {self.data_path}")

	# Pre-compute token lengths to filter out examples that exceed sequence_length.
	# This is fast (just encoding, no tensor creation) and avoids wasting time
	# on examples that would be skipped anyway.
	valid_indices: list[int] = []
	skipped = 0
	for i, ex in enumerate(self.raw_examples):
	token_len = self._estimate_token_length(ex)
	if token_len <= self.sequence_length and token_len >= 10:
	valid_indices.append(i)
	else:
	skipped += 1

	if not valid_indices:
	raise ValueError(
	f"All examples were skipped (too long for sequence_length={sequence_length}). "
	f"Consider increasing sequence_length."
	)

	if skipped > 0:
	print(
	f" InstructionDataset: skipped {skipped} examples "
	f"(exceeded sequence_length={sequence_length})"
	)

	self.valid_indices = valid_indices

	def _render_turn(self, role: str, content: str) -> str:
	"""Render a single conversation turn with its role marker."""
	if role == "system":
	return f"{self.SYSTEM_MARKER}\n{content}\n"
	elif role == "user":
	return f"{self.USER_MARKER}\n{content}\n"
	elif role == "assistant":
	return f"{self.ASSISTANT_MARKER}\n{content}</s>\n"
	else:
	raise ValueError(f"Unknown role: {role}")

	def _estimate_token_length(self, ex: dict) -> int:
	"""Exact token count for filtering.

	Mirrors the per-turn tokenization of _process_example exactly: BPE is
	not additive across turn boundaries, so encoding the concatenated text
	in one pass can undercount versus per-turn encoding.
	"""
	conversations = ex.get("conversations", [])
	if not conversations or len(conversations) < 2:
	return 0
	if conversations[-1]["role"] != "assistant":
	return 0
	total = 0
	for turn in conversations:
	role = turn.get("role", "")
	content = (turn.get("content") or "").strip()
	if not content:
	continue
	if role in ("system", "user"):
	text = self._render_turn(role, content)
	total += len(self.tokenizer.encode(text).ids)
	elif role == "assistant":
	marker_text = f"{self.ASSISTANT_MARKER}\n"
	total += len(self.tokenizer.encode(marker_text).ids)
	content_text = f"{content}</s>\n"
	total += len(self.tokenizer.encode(content_text).ids)
	return total

	def _process_example(self, ex: dict) -> dict[str, torch.Tensor] \| None:
	"""Tokenize a conversation example and build labels with masking.

	Returns None if the example is too long for sequence_length.
	"""
	conversations = ex.get("conversations", [])
	if not conversations or len(conversations) < 2:
	return None

	# Validate: must end with assistant turn.
	if conversations[-1]["role"] != "assistant":
	return None

	input_ids: list[int] = []
	labels: list[int] = [] # -100 for masked, token_id for assistant content + </s>

	for turn in conversations:
	role = turn.get("role", "")
	content = (turn.get("content") or "").strip()
	if not content:
	continue

	if role == "system":
	# Render and tokenize the marker + content. All masked.
	text = self._render_turn(role, content)
	ids = self.tokenizer.encode(text).ids
	input_ids.extend(ids)
	labels.extend([-100] * len(ids))

	elif role == "user":
	# Render and tokenize the marker + content. All masked.
	text = self._render_turn(role, content)
	ids = self.tokenizer.encode(text).ids
	input_ids.extend(ids)
	labels.extend([-100] * len(ids))

	elif role == "assistant":
	# Render marker (masked) + content + </s> (not masked).
	marker_text = f"{self.ASSISTANT_MARKER}\n"
	marker_ids = self.tokenizer.encode(marker_text).ids

	content_text = f"{content}</s>\n"
	content_ids = self.tokenizer.encode(content_text).ids

	input_ids.extend(marker_ids)
	labels.extend([-100] * len(marker_ids))

	input_ids.extend(content_ids)
	labels.extend(content_ids)

	else:
	# Unknown role, skip.
	continue

	# Check length: skip if too long (don't truncate - would lose </s>).
	if len(input_ids) > self.sequence_length:
	return None
	if len(input_ids) < 10:
	return None

	actual_length = len(input_ids)
	padding_length = self.sequence_length - actual_length

	# Pad input_ids with pad_token_id.
	input_ids_padded = input_ids + [self.pad_token_id] * padding_length

	# Pad labels with -100 (ignored in loss).
	labels_padded = labels + [-100] * padding_length

	# Attention mask: 1 for real tokens, 0 for padding.
	attention_mask = [1] * actual_length + [0] * padding_length

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

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

	def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
	if idx < 0 or idx >= len(self.valid_indices):
	raise IndexError(
	f"Index {idx} out of range for dataset with {len(self.valid_indices)} examples"
	)
	raw_idx = self.valid_indices[idx]
	ex = self.raw_examples[raw_idx]
	processed = self._process_example(ex)
	# processed should not be None since we pre-filtered, but guard anyway.
	if processed is None:
	# Fallback: return a minimal valid example (shouldn't happen).
	raise RuntimeError(
	f"Example at index {idx} (raw {raw_idx}) failed to process "
	f"despite passing pre-filter. This is a bug."
	)
	return processed


	if __name__ == "__main__":
	import sys

	print("BoccaccioAI Dataset classes")
	print(f" PreTokenizedDataset: for pre-training on binary memmap files")
	print(f" InstructionDataset: for instruction fine-tuning on JSONL data")

	if len(sys.argv) > 1 and sys.argv[1] == "--test-pretokenized":
	if len(sys.argv) < 3:
	print("Usage: python dataset.py --test-pretokenized <data_path>")
	sys.exit(1)
	ds = PreTokenizedDataset(sys.argv[2])
	print(f" Loaded {len(ds)} sequences from {sys.argv[2]}")
	sample = ds[0]
	print(f" Sample input_ids shape: {sample['input_ids'].shape}")
	print(f" Sample input_ids dtype: {sample['input_ids'].dtype}")
	print(f" First 20 token IDs: {sample['input_ids'][:20].tolist()}")

	if len(sys.argv) > 1 and sys.argv[1] == "--test-instruction":
	if len(sys.argv) < 4:
	print("Usage: python dataset.py --test-instruction <data_path> <tokenizer_path>")
	sys.exit(1)
	ds = InstructionDataset(sys.argv[2], sys.argv[3])
	print(f" Loaded {len(ds)} examples from {sys.argv[2]}")
	sample = ds[0]
	print(f" Sample input_ids shape: {sample['input_ids'].shape}")
	print(f" Sample labels shape: {sample['labels'].shape}")
	print(f" Sample attn_mask shape: {sample['attention_mask'].shape}")
	num_masked = (sample["labels"] == -100).sum().item()
	num_active = sample["labels"].numel() - num_masked
	print(f" Masked label positions: {num_masked}")
	print(f" Active label positions: {num_active} (assistant tokens + </s>)")