distil-trainer 0.1.10__py3-none-any.whl

distil_trainer/data/__init__.py

@@ -0,0 +1,19 @@
+"""Data module for distillation training."""
+
+from distil_trainer.data.datamodule import DistillationDataModule
+from distil_trainer.data.datasets import (
+    SentenceDistillationDataset,
+    TripletDataset,
+    ParallelSentencesDataset,
+)
+from distil_trainer.data.collators import DistillationCollator
+from distil_trainer.data.loaders import DatasetLoaders
+
+__all__ = [
+    "DistillationDataModule",
+    "SentenceDistillationDataset",
+    "TripletDataset",
+    "ParallelSentencesDataset",
+    "DistillationCollator",
+    "DatasetLoaders",
+]

distil_trainer/data/collators.py

@@ -0,0 +1,240 @@
+"""Data collators for distillation training."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+class DistillationCollator:
+    """
+    Collator for distillation training data.
+
+    Handles tokenization and batching of text data.
+
+    Example:
+        >>> collator = DistillationCollator(tokenizer, max_length=512)
+        >>> batch = collator([{"sentence": "Hello"}, {"sentence": "World"}])
+    """
+
+    def __init__(
+        self,
+        tokenizer: Any = None,
+        max_length: int = 512,
+        text_column: str = "sentence",
+        padding: bool = True,
+        truncation: bool = True,
+    ):
+        """
+        Initialize the collator.
+
+        Args:
+            tokenizer: Tokenizer for encoding text.
+            max_length: Maximum sequence length.
+            text_column: Name of the text column.
+            padding: Whether to pad sequences.
+            truncation: Whether to truncate sequences.
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.text_column = text_column
+        self.padding = padding
+        self.truncation = truncation
+
+    def __call__(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+        """
+        Collate a batch of samples.
+
+        Args:
+            batch: List of sample dictionaries.
+
+        Returns:
+            Collated batch dictionary.
+        """
+        # Extract text
+        texts = [sample.get(self.text_column, sample.get("sentence", "")) for sample in batch]
+
+        # Tokenize if tokenizer is available
+        if self.tokenizer is not None:
+            encoded = self.tokenizer(
+                texts,
+                padding=self.padding,
+                truncation=self.truncation,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            result = dict(encoded)
+        else:
+            result = {self.text_column: texts}
+
+        # Add labels if present
+        if "label" in batch[0]:
+            labels = [sample["label"] for sample in batch]
+            if isinstance(labels[0], torch.Tensor):
+                result["label"] = torch.stack(labels)
+            else:
+                result["label"] = torch.tensor(labels)
+
+        return result
+
+
+class TripletCollator:
+    """
+    Collator for triplet data (query, positive, negatives).
+
+    Example:
+        >>> collator = TripletCollator(tokenizer)
+        >>> batch = collator([{"query": "q1", "positive": "p1", "negatives": ["n1"]}])
+    """
+
+    def __init__(
+        self,
+        tokenizer: Any = None,
+        max_length: int = 512,
+        padding: bool = True,
+        truncation: bool = True,
+    ):
+        """
+        Initialize the collator.
+
+        Args:
+            tokenizer: Tokenizer for encoding text.
+            max_length: Maximum sequence length.
+            padding: Whether to pad sequences.
+            truncation: Whether to truncate sequences.
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.padding = padding
+        self.truncation = truncation
+
+    def __call__(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+        """Collate a batch of triplet samples."""
+        queries = [sample["query"] for sample in batch]
+        positives = [sample["positive"] for sample in batch]
+
+        result = {}
+
+        if self.tokenizer is not None:
+            # Tokenize queries
+            query_encoded = self.tokenizer(
+                queries,
+                padding=self.padding,
+                truncation=self.truncation,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            result["query_input_ids"] = query_encoded["input_ids"]
+            result["query_attention_mask"] = query_encoded["attention_mask"]
+
+            # Tokenize positives
+            pos_encoded = self.tokenizer(
+                positives,
+                padding=self.padding,
+                truncation=self.truncation,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            result["positive_input_ids"] = pos_encoded["input_ids"]
+            result["positive_attention_mask"] = pos_encoded["attention_mask"]
+
+            # Tokenize negatives if present
+            if "negatives" in batch[0]:
+                all_negatives = []
+                for sample in batch:
+                    all_negatives.extend(sample.get("negatives", []))
+
+                if all_negatives:
+                    neg_encoded = self.tokenizer(
+                        all_negatives,
+                        padding=self.padding,
+                        truncation=self.truncation,
+                        max_length=self.max_length,
+                        return_tensors="pt",
+                    )
+                    result["negative_input_ids"] = neg_encoded["input_ids"]
+                    result["negative_attention_mask"] = neg_encoded["attention_mask"]
+
+        else:
+            result["query"] = queries
+            result["positive"] = positives
+            if "negatives" in batch[0]:
+                result["negatives"] = [sample.get("negatives", []) for sample in batch]
+
+        return result
+
+
+class ParallelSentenceCollator:
+    """
+    Collator for parallel sentence data (multilingual).
+
+    Example:
+        >>> collator = ParallelSentenceCollator(tokenizer)
+        >>> batch = collator([{"english": "Hello", "non_english": "Hallo"}])
+    """
+
+    def __init__(
+        self,
+        tokenizer: Any = None,
+        max_length: int = 128,
+        padding: bool = True,
+        truncation: bool = True,
+    ):
+        """
+        Initialize the collator.
+
+        Args:
+            tokenizer: Tokenizer for encoding text.
+            max_length: Maximum sequence length.
+            padding: Whether to pad sequences.
+            truncation: Whether to truncate sequences.
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.padding = padding
+        self.truncation = truncation
+
+    def __call__(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+        """Collate a batch of parallel sentence samples."""
+        english = [sample["english"] for sample in batch]
+        non_english = [sample["non_english"] for sample in batch]
+
+        result = {}
+
+        if self.tokenizer is not None:
+            # Tokenize English
+            en_encoded = self.tokenizer(
+                english,
+                padding=self.padding,
+                truncation=self.truncation,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            result["english_input_ids"] = en_encoded["input_ids"]
+            result["english_attention_mask"] = en_encoded["attention_mask"]
+
+            # Tokenize non-English
+            ne_encoded = self.tokenizer(
+                non_english,
+                padding=self.padding,
+                truncation=self.truncation,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            result["non_english_input_ids"] = ne_encoded["input_ids"]
+            result["non_english_attention_mask"] = ne_encoded["attention_mask"]
+
+        else:
+            result["english"] = english
+            result["non_english"] = non_english
+
+        # Add labels if present
+        if "label" in batch[0]:
+            labels = [sample["label"] for sample in batch]
+            if isinstance(labels[0], torch.Tensor):
+                result["label"] = torch.stack(labels)
+            else:
+                result["label"] = torch.tensor(labels)
+
+        return result

distil_trainer/data/datamodule.py

@@ -0,0 +1,191 @@
+"""Base data module for distillation training."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+from datasets import Dataset, DatasetDict, load_dataset
+from torch.utils.data import DataLoader
+
+from distil_trainer.data.collators import DistillationCollator
+
+logger = logging.getLogger(__name__)
+
+
+class DistillationDataModule:
+    """
+    Base data module for distillation training.
+
+    Handles data loading, preprocessing, and DataLoader creation.
+
+    Example:
+        >>> datamodule = DistillationDataModule(
+        ...     train_data="sentence-transformers/all-nli",
+        ...     text_column="sentence",
+        ...     batch_size=32,
+        ... )
+        >>> datamodule.prepare_data()
+        >>> datamodule.setup()
+        >>> train_loader = datamodule.train_dataloader()
+    """
+
+    def __init__(
+        self,
+        train_data: str | Dataset | None = None,
+        eval_data: str | Dataset | None = None,
+        test_data: str | Dataset | None = None,
+        tokenizer: Any = None,
+        text_column: str = "sentence",
+        max_seq_length: int = 512,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        preprocessing_fn: Callable | None = None,
+        max_samples: int | None = None,
+    ):
+        """
+        Initialize the data module.
+
+        Args:
+            train_data: Training data path/name or Dataset.
+            eval_data: Evaluation data path/name or Dataset.
+            test_data: Test data path/name or Dataset.
+            tokenizer: Tokenizer for encoding text.
+            text_column: Name of the text column in the dataset.
+            max_seq_length: Maximum sequence length for tokenization.
+            batch_size: Batch size for DataLoaders.
+            num_workers: Number of workers for data loading.
+            preprocessing_fn: Optional preprocessing function.
+            max_samples: Maximum number of samples to use.
+        """
+        self.train_data = train_data
+        self.eval_data = eval_data
+        self.test_data = test_data
+        self.tokenizer = tokenizer
+        self.text_column = text_column
+        self.max_seq_length = max_seq_length
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.preprocessing_fn = preprocessing_fn
+        self.max_samples = max_samples
+
+        # Datasets
+        self.train_dataset: Dataset | None = None
+        self.eval_dataset: Dataset | None = None
+        self.test_dataset: Dataset | None = None
+
+        # Collator
+        self.collator = DistillationCollator(
+            tokenizer=tokenizer,
+            max_length=max_seq_length,
+            text_column=text_column,
+        )
+
+    def prepare_data(self) -> None:
+        """Download and prepare data. Called once on a single process."""
+        # Download datasets if they're specified as strings
+        if isinstance(self.train_data, str):
+            logger.info(f"Downloading training data: {self.train_data}")
+            _ = load_dataset(self.train_data, split="train")
+
+        if isinstance(self.eval_data, str):
+            logger.info(f"Downloading eval data: {self.eval_data}")
+            _ = load_dataset(self.eval_data, split="validation")
+
+    def setup(self, stage: str | None = None) -> None:
+        """
+        Set up datasets for training/evaluation.
+
+        Args:
+            stage: "fit", "validate", "test", or None for all.
+        """
+        if stage in (None, "fit"):
+            self.train_dataset = self._load_dataset(self.train_data, "train")
+            if self.eval_data is not None:
+                self.eval_dataset = self._load_dataset(self.eval_data, "validation")
+
+        if stage in (None, "validate"):
+            if self.eval_dataset is None and self.eval_data is not None:
+                self.eval_dataset = self._load_dataset(self.eval_data, "validation")
+
+        if stage in (None, "test"):
+            if self.test_data is not None:
+                self.test_dataset = self._load_dataset(self.test_data, "test")
+
+    def _load_dataset(self, data: str | Dataset, split: str = "train") -> Dataset:
+        """Load and preprocess a dataset."""
+        if data is None:
+            return None
+
+        if isinstance(data, Dataset):
+            dataset = data
+        else:
+            try:
+                dataset = load_dataset(data, split=split)
+            except Exception:
+                # Try loading without split
+                loaded = load_dataset(data)
+                if isinstance(loaded, DatasetDict):
+                    if split in loaded:
+                        dataset = loaded[split]
+                    else:
+                        dataset = list(loaded.values())[0]
+                else:
+                    dataset = loaded
+
+        # Limit samples if specified
+        if self.max_samples is not None and len(dataset) > self.max_samples:
+            dataset = dataset.select(range(self.max_samples))
+
+        # Apply preprocessing
+        if self.preprocessing_fn is not None:
+            dataset = dataset.map(
+                self.preprocessing_fn,
+                batched=True,
+                remove_columns=dataset.column_names,
+            )
+
+        logger.info(f"Loaded dataset with {len(dataset)} samples")
+        return dataset
+
+    def train_dataloader(self) -> DataLoader:
+        """Return training DataLoader."""
+        if self.train_dataset is None:
+            raise ValueError("Training dataset not loaded. Call setup() first.")
+
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            collate_fn=self.collator,
+            num_workers=self.num_workers,
+            pin_memory=True,
+        )
+
+    def val_dataloader(self) -> DataLoader | None:
+        """Return validation DataLoader."""
+        if self.eval_dataset is None:
+            return None
+
+        return DataLoader(
+            self.eval_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            collate_fn=self.collator,
+            num_workers=self.num_workers,
+            pin_memory=True,
+        )
+
+    def test_dataloader(self) -> DataLoader | None:
+        """Return test DataLoader."""
+        if self.test_dataset is None:
+            return None
+
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            collate_fn=self.collator,
+            num_workers=self.num_workers,
+            pin_memory=True,
+        )

distil_trainer/data/datasets.py

@@ -0,0 +1,245 @@
+"""Dataset classes for distillation training."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch.utils.data import Dataset as TorchDataset
+
+
+class SentenceDistillationDataset(TorchDataset):
+    """
+    Dataset for sentence embedding distillation.
+
+    Example:
+        >>> dataset = SentenceDistillationDataset(
+        ...     sentences=["Hello world", "How are you?"],
+        ...     teacher_embeddings=teacher_embeddings,
+        ...     tokenizer=tokenizer,
+        ... )
+        >>> item = dataset[0]
+    """
+
+    def __init__(
+        self,
+        sentences: list[str],
+        teacher_embeddings: torch.Tensor | list | None = None,
+        tokenizer: Any = None,
+        max_length: int = 512,
+    ):
+        """
+        Initialize the dataset.
+
+        Args:
+            sentences: List of sentences to encode.
+            teacher_embeddings: Precomputed teacher embeddings.
+            tokenizer: Tokenizer for encoding text.
+            max_length: Maximum sequence length.
+        """
+        self.sentences = sentences
+        self.teacher_embeddings = teacher_embeddings
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self) -> int:
+        return len(self.sentences)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        item = {"sentence": self.sentences[idx]}
+
+        if self.teacher_embeddings is not None:
+            if isinstance(self.teacher_embeddings, torch.Tensor):
+                item["label"] = self.teacher_embeddings[idx]
+            else:
+                item["label"] = torch.tensor(self.teacher_embeddings[idx])
+
+        return item
+
+
+class TripletDataset(TorchDataset):
+    """
+    Dataset for triplet (query, positive, negative) training.
+
+    Used for contrastive learning and retrieval model training.
+
+    Example:
+        >>> dataset = TripletDataset(
+        ...     queries=["What is Python?"],
+        ...     positive_docs=[["Python is a programming language..."]],
+        ...     negative_docs=[["Java is a programming language..."]]
+        ... )
+    """
+
+    def __init__(
+        self,
+        queries: list[str],
+        positive_docs: list[list[str]],
+        negative_docs: list[list[str]] | None = None,
+        tokenizer: Any = None,
+        max_length: int = 512,
+        num_negatives: int = 5,
+    ):
+        """
+        Initialize the dataset.
+
+        Args:
+            queries: List of query strings.
+            positive_docs: List of positive document lists (multiple per query).
+            negative_docs: List of negative document lists.
+            tokenizer: Tokenizer for encoding text.
+            max_length: Maximum sequence length.
+            num_negatives: Number of negatives to sample per query.
+        """
+        self.queries = queries
+        self.positive_docs = positive_docs
+        self.negative_docs = negative_docs
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.num_negatives = num_negatives
+
+    def __len__(self) -> int:
+        return len(self.queries)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        query = self.queries[idx]
+        positives = self.positive_docs[idx]
+
+        # Select first positive (or random)
+        positive = positives[0] if positives else ""
+
+        item = {
+            "query": query,
+            "positive": positive,
+        }
+
+        # Add negatives if available
+        if self.negative_docs is not None and idx < len(self.negative_docs):
+            negatives = self.negative_docs[idx]
+            # Limit to num_negatives
+            item["negatives"] = negatives[: self.num_negatives]
+
+        return item
+
+
+class ParallelSentencesDataset(TorchDataset):
+    """
+    Dataset for parallel sentences (multilingual training).
+
+    Example:
+        >>> dataset = ParallelSentencesDataset(
+        ...     source_sentences=["Hello"],
+        ...     target_sentences=["Hallo"],
+        ...     teacher_embeddings=embeddings,
+        ... )
+    """
+
+    def __init__(
+        self,
+        source_sentences: list[str],
+        target_sentences: list[str],
+        teacher_embeddings: torch.Tensor | list | None = None,
+        tokenizer: Any = None,
+        max_length: int = 128,
+    ):
+        """
+        Initialize the dataset.
+
+        Args:
+            source_sentences: Source language sentences.
+            target_sentences: Target language sentences.
+            teacher_embeddings: Teacher embeddings for source sentences.
+            tokenizer: Tokenizer for encoding text.
+            max_length: Maximum sequence length.
+        """
+        assert len(source_sentences) == len(target_sentences)
+
+        self.source_sentences = source_sentences
+        self.target_sentences = target_sentences
+        self.teacher_embeddings = teacher_embeddings
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self) -> int:
+        return len(self.source_sentences)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        item = {
+            "english": self.source_sentences[idx],
+            "non_english": self.target_sentences[idx],
+        }
+
+        if self.teacher_embeddings is not None:
+            if isinstance(self.teacher_embeddings, torch.Tensor):
+                item["label"] = self.teacher_embeddings[idx]
+            else:
+                item["label"] = torch.tensor(self.teacher_embeddings[idx])
+
+        return item
+
+
+class ReasoningDataset(TorchDataset):
+    """
+    Dataset for chain-of-thought/reasoning distillation.
+
+    Example:
+        >>> dataset = ReasoningDataset(
+        ...     questions=["What is 2+2?"],
+        ...     reasoning_chains=["Let me think... 2+2=4"],
+        ...     answers=["4"],
+        ... )
+    """
+
+    def __init__(
+        self,
+        questions: list[str],
+        reasoning_chains: list[str],
+        answers: list[str],
+        system_prompts: list[str] | None = None,
+        tokenizer: Any = None,
+        max_length: int = 16384,
+        chat_format: bool = True,
+    ):
+        """
+        Initialize the dataset.
+
+        Args:
+            questions: List of questions.
+            reasoning_chains: List of reasoning chains.
+            answers: List of final answers.
+            system_prompts: Optional system prompts.
+            tokenizer: Tokenizer for encoding.
+            max_length: Maximum sequence length.
+            chat_format: Whether to format as chat.
+        """
+        self.questions = questions
+        self.reasoning_chains = reasoning_chains
+        self.answers = answers
+        self.system_prompts = system_prompts
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.chat_format = chat_format
+
+    def __len__(self) -> int:
+        return len(self.questions)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        question = self.questions[idx]
+        reasoning = self.reasoning_chains[idx]
+        answer = self.answers[idx]
+
+        item = {
+            "question": question,
+            "reasoning": reasoning,
+            "answer": answer,
+        }
+
+        if self.system_prompts is not None and idx < len(self.system_prompts):
+            item["system"] = self.system_prompts[idx]
+
+        if self.chat_format:
+            # Format as conversation
+            full_response = f"<|begin_of_thought|>\n{reasoning}\n<|end_of_thought|>\n\n<|begin_of_solution|>\n{answer}\n<|end_of_solution|>"
+            item["full_response"] = full_response
+
+        return item