bead 0.1.0__py3-none-any.whl

bead/active_learning/trainers/huggingface.py

@@ -0,0 +1,304 @@
+"""HuggingFace Transformers trainer implementation.
+
+This module provides a trainer that uses the HuggingFace Transformers library
+for model training with integrated TensorBoard logging.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from bead.active_learning.trainers.base import BaseTrainer, ModelMetadata
+from bead.data.base import BeadBaseModel
+from bead.data.timestamps import format_iso8601, now_iso8601
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+
+
+class HuggingFaceTrainer(BaseTrainer):
+    """Trainer using HuggingFace Transformers.
+
+    This trainer uses the HuggingFace Transformers library to train models
+    for sequence classification and other NLP tasks. It supports TensorBoard
+    logging and checkpoint management.
+
+    Parameters
+    ----------
+    config : dict[str, int | str | float | bool | Path] | BeadBaseModel
+        Training configuration with the following expected fields:
+        - model_name: str - Base model name/path
+        - task_type: str - Task type (classification, regression, etc.)
+        - num_labels: int | None - Number of labels for classification
+        - output_dir: Path - Directory for outputs
+        - num_epochs: int - Number of training epochs
+        - batch_size: int - Training batch size
+        - learning_rate: float - Learning rate
+        - weight_decay: float - Weight decay
+        - warmup_steps: int - Warmup steps
+        - evaluation_strategy: str - Evaluation strategy (epoch, steps, no)
+        - save_strategy: str - Save strategy (epoch, steps, no)
+        - load_best_model_at_end: bool - Load best model at end
+        - logging_dir: Path | None - Logging directory
+        - fp16: bool - Use mixed precision
+
+    Attributes
+    ----------
+    config : dict[str, int | str | float | bool | Path] | BeadBaseModel
+        Training configuration.
+    model : PreTrainedModel | None
+        The trained model.
+    tokenizer : PreTrainedTokenizer | None
+        The tokenizer.
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> config = {
+    ...     "model_name": "bert-base-uncased",
+    ...     "task_type": "classification",
+    ...     "num_labels": 2,
+    ...     "output_dir": Path("output"),
+    ...     "num_epochs": 3,
+    ...     "batch_size": 16,
+    ...     "learning_rate": 2e-5,
+    ...     "weight_decay": 0.01,
+    ...     "warmup_steps": 0,
+    ...     "evaluation_strategy": "epoch",
+    ...     "save_strategy": "epoch",
+    ...     "load_best_model_at_end": True,
+    ...     "logging_dir": None,
+    ...     "fp16": False
+    ... }
+    >>> trainer = HuggingFaceTrainer(config)
+    >>> trainer.model is None
+    True
+    """
+
+    def __init__(
+        self, config: dict[str, int | str | float | bool | Path] | BeadBaseModel
+    ) -> None:
+        super().__init__(config)
+        self.model: PreTrainedModel | None = None
+        self.tokenizer: PreTrainedTokenizer | None = None
+
+    def _get_config_value(
+        self, key: str, default: int | str | float | bool | Path | None = None
+    ) -> int | str | float | bool | Path | None:
+        """Get configuration value with fallback to default.
+
+        Parameters
+        ----------
+        key : str
+            Configuration key.
+        default : int | str | float | bool | Path | None
+            Default value if key not found.
+
+        Returns
+        -------
+        int | str | float | bool | Path | None
+            Configuration value.
+        """
+        if hasattr(self.config, key):
+            return getattr(self.config, key)
+        if isinstance(self.config, dict):
+            return self.config.get(key, default)
+        return default
+
+    def train(
+        self, train_data: Dataset, eval_data: Dataset | None = None
+    ) -> ModelMetadata:
+        """Train model using HuggingFace Trainer.
+
+        Parameters
+        ----------
+        train_data : Dataset
+            HuggingFace Dataset for training.
+        eval_data : Dataset | None
+            HuggingFace Dataset for evaluation.
+
+        Returns
+        -------
+        ModelMetadata
+            Training metadata.
+
+        Raises
+        ------
+        ValueError
+            If task type is not supported.
+
+        Examples
+        --------
+        >>> config = {"model_name": "bert-base-uncased"}  # doctest: +SKIP
+        >>> trainer = HuggingFaceTrainer(config)  # doctest: +SKIP
+        >>> metadata = trainer.train(train_dataset)  # doctest: +SKIP
+        >>> metadata.framework  # doctest: +SKIP
+        'huggingface'
+        """
+        from transformers import (  # noqa: PLC0415
+            AutoModelForSequenceClassification,
+            AutoTokenizer,
+            DataCollatorWithPadding,
+            Trainer,
+            TrainingArguments,
+        )
+
+        start_time = time.time()
+
+        # Get config values
+        model_name = self._get_config_value("model_name", "bert-base-uncased")
+        task_type = self._get_config_value("task_type", "classification")
+        num_labels = self._get_config_value("num_labels", 2)
+        output_dir = self._get_config_value("output_dir", Path("output"))
+        num_epochs = self._get_config_value("num_epochs", 3)
+        batch_size = self._get_config_value("batch_size", 16)
+        learning_rate = self._get_config_value("learning_rate", 2e-5)
+        weight_decay = self._get_config_value("weight_decay", 0.01)
+        warmup_steps = self._get_config_value("warmup_steps", 0)
+        evaluation_strategy = self._get_config_value("evaluation_strategy", "epoch")
+        save_strategy = self._get_config_value("save_strategy", "epoch")
+        load_best = self._get_config_value("load_best_model_at_end", True)
+        logging_dir = self._get_config_value("logging_dir", None)
+        fp16 = self._get_config_value("fp16", False)
+
+        # Load model and tokenizer
+        if task_type == "classification":
+            self.model = AutoModelForSequenceClassification.from_pretrained(
+                model_name, num_labels=num_labels
+            )
+        else:
+            msg = f"Task type not supported: {task_type}"
+            raise ValueError(msg)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+        # Create training arguments
+        training_args = TrainingArguments(
+            output_dir=str(output_dir),
+            num_train_epochs=num_epochs,
+            per_device_train_batch_size=batch_size,
+            per_device_eval_batch_size=batch_size,
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            warmup_steps=warmup_steps,
+            eval_strategy=evaluation_strategy,  # type: ignore
+            save_strategy=save_strategy,
+            load_best_model_at_end=load_best,
+            logging_dir=str(logging_dir) if logging_dir else None,
+            fp16=fp16,
+            report_to=["tensorboard"] if logging_dir else [],
+        )
+
+        # Create data collator
+        data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
+
+        # Create trainer
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_data,
+            eval_dataset=eval_data,
+            data_collator=data_collator,
+        )
+
+        # Train
+        trainer.train()
+
+        # Evaluate
+        metrics = {}
+        if eval_data is not None:
+            eval_results = trainer.evaluate()
+            metrics = {k: float(v) for k, v in eval_results.items()}
+
+        training_time = time.time() - start_time
+
+        # Get best checkpoint path
+        best_checkpoint = None
+        if trainer.state.best_model_checkpoint:
+            best_checkpoint = Path(trainer.state.best_model_checkpoint)
+
+        # Create metadata
+        config_dict = (
+            self.config
+            if isinstance(self.config, dict)
+            else (
+                self.config.model_dump() if hasattr(self.config, "model_dump") else {}
+            )
+        )
+
+        metadata = ModelMetadata(
+            model_name=model_name,
+            framework="huggingface",
+            training_config=config_dict,
+            training_data_path=Path("train.json"),
+            eval_data_path=Path("eval.json") if eval_data else None,
+            metrics=metrics,
+            best_checkpoint=best_checkpoint,
+            training_time=training_time,
+            training_timestamp=format_iso8601(now_iso8601()),
+        )
+
+        return metadata
+
+    def save_model(self, output_dir: Path, metadata: ModelMetadata) -> None:
+        """Save model and metadata.
+
+        Parameters
+        ----------
+        output_dir : Path
+            Directory to save model and metadata.
+        metadata : ModelMetadata
+            Training metadata to save.
+
+        Examples
+        --------
+        >>> trainer = HuggingFaceTrainer({})  # doctest: +SKIP
+        >>> trainer.save_model(Path("output"), metadata)  # doctest: +SKIP
+        """
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Save model
+        if self.model is not None:
+            self.model.save_pretrained(output_dir / "model")
+        if self.tokenizer is not None:
+            self.tokenizer.save_pretrained(output_dir / "model")
+
+        # Save metadata
+        with open(output_dir / "metadata.json", "w") as f:
+            # Convert Path objects to strings for JSON serialization
+            metadata_dict = metadata.model_dump()
+            json.dump(metadata_dict, f, indent=2, default=str)
+
+    def load_model(self, model_dir: Path) -> PreTrainedModel:
+        """Load model.
+
+        Parameters
+        ----------
+        model_dir : Path
+            Directory containing saved model.
+
+        Returns
+        -------
+        PreTrainedModel
+            Loaded model.
+
+        Examples
+        --------
+        >>> trainer = HuggingFaceTrainer({})  # doctest: +SKIP
+        >>> model = trainer.load_model(Path("saved_model"))  # doctest: +SKIP
+        """
+        from transformers import (  # noqa: PLC0415
+            AutoModelForSequenceClassification,
+            AutoTokenizer,
+        )
+
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            model_dir / "model"
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir / "model")
+
+        return self.model

bead/active_learning/trainers/lightning.py

@@ -0,0 +1,324 @@
+"""PyTorch Lightning trainer implementation.
+
+This module provides a trainer that uses PyTorch Lightning for model training
+with callbacks for checkpointing and early stopping.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from bead.active_learning.trainers.base import BaseTrainer, ModelMetadata
+from bead.data.base import BeadBaseModel
+from bead.data.timestamps import format_iso8601, now_iso8601
+
+if TYPE_CHECKING:
+    import pytorch_lightning as pl
+    from torch.nn import Module
+    from torch.utils.data import DataLoader
+
+
+def create_lightning_module(
+    model: Module, learning_rate: float = 2e-5
+) -> pl.LightningModule:
+    """Create a PyTorch Lightning module.
+
+    Parameters
+    ----------
+    model
+        The PyTorch model to wrap in a Lightning module.
+    learning_rate
+        Learning rate for the AdamW optimizer.
+
+    Returns
+    -------
+    pl.LightningModule
+        Lightning module wrapping the provided model with training,
+        validation, and optimizer configuration.
+    """
+    import pytorch_lightning as pl  # noqa: PLC0415
+    import torch  # noqa: PLC0415
+
+    class _LightningModule(pl.LightningModule):
+        def __init__(self) -> None:
+            super().__init__()
+            self.model = model
+            self.learning_rate = learning_rate
+
+        def forward(self, **inputs: Any) -> Any:
+            return self.model(**inputs)
+
+        def training_step(self, batch: Any, batch_idx: int) -> Any:
+            outputs = self(**batch)
+            loss = outputs.loss
+            self.log("train_loss", loss)
+            return loss
+
+        def validation_step(self, batch: Any, batch_idx: int) -> Any:
+            outputs = self(**batch)
+            loss = outputs.loss
+            self.log("val_loss", loss)
+            return loss
+
+        def configure_optimizers(self) -> Any:
+            optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
+            return optimizer
+
+    return _LightningModule()
+
+
+class PyTorchLightningTrainer(BaseTrainer):
+    """Trainer using PyTorch Lightning.
+
+    Trains models using PyTorch Lightning with callbacks for checkpointing
+    and early stopping.
+
+    Parameters
+    ----------
+    config
+        Training configuration as a dict or config object with the following
+        fields:
+
+        - model_name: str, base model name or path
+        - num_labels: int, number of output labels
+        - num_epochs: int, number of training epochs
+        - learning_rate: float, learning rate for optimizer
+        - output_dir: Path, directory for outputs and checkpoints
+        - logging_dir: Path or None, optional TensorBoard logging directory
+
+    Attributes
+    ----------
+    config : dict[str, int | str | float | bool | Path] | BeadBaseModel
+        Training configuration.
+    lightning_module : pl.LightningModule | None
+        The Lightning module wrapper, set after training.
+
+    Examples
+    --------
+    >>> from pathlib import Path
+    >>> config = {
+    ...     "model_name": "bert-base-uncased",
+    ...     "num_labels": 2,
+    ...     "num_epochs": 3,
+    ...     "learning_rate": 2e-5,
+    ...     "output_dir": Path("output"),
+    ...     "logging_dir": None
+    ... }
+    >>> trainer = PyTorchLightningTrainer(config)
+    >>> trainer.lightning_module is None
+    True
+    """
+
+    def __init__(
+        self, config: dict[str, int | str | float | bool | Path] | BeadBaseModel
+    ) -> None:
+        super().__init__(config)
+        self.lightning_module: pl.LightningModule | None = None
+
+    def _get_config_value(
+        self, key: str, default: int | str | float | bool | Path | None = None
+    ) -> int | str | float | bool | Path | None:
+        """Get configuration value with fallback to default.
+
+        Parameters
+        ----------
+        key
+            Configuration key to retrieve.
+        default
+            Default value if key is not found.
+
+        Returns
+        -------
+        int | str | float | bool | Path | None
+            Configuration value for the given key, or default if not found.
+        """
+        if hasattr(self.config, key):
+            return getattr(self.config, key)
+        if isinstance(self.config, dict):
+            return self.config.get(key, default)
+        return default
+
+    def train(
+        self, train_data: DataLoader, eval_data: DataLoader | None = None
+    ) -> ModelMetadata:
+        """Train a model using PyTorch Lightning.
+
+        Loads a pretrained model, wraps it in a Lightning module, and trains
+        with checkpointing and early stopping callbacks.
+
+        Parameters
+        ----------
+        train_data
+            Training dataloader providing batches for training.
+        eval_data
+            Optional evaluation dataloader for validation during training.
+
+        Returns
+        -------
+        ModelMetadata
+            Metadata containing model name, framework, training config,
+            metrics, checkpoint path, and training time.
+
+        Examples
+        --------
+        >>> config = {"model_name": "bert-base-uncased"}  # doctest: +SKIP
+        >>> trainer = PyTorchLightningTrainer(config)  # doctest: +SKIP
+        >>> metadata = trainer.train(train_loader)  # doctest: +SKIP
+        >>> metadata.framework  # doctest: +SKIP
+        'pytorch_lightning'
+        """
+        import pytorch_lightning as pl  # noqa: PLC0415
+        from transformers import AutoModelForSequenceClassification  # noqa: PLC0415
+
+        start_time = time.time()
+
+        # get config values
+        model_name = self._get_config_value("model_name", "bert-base-uncased")
+        num_labels = self._get_config_value("num_labels", 2)
+        num_epochs = self._get_config_value("num_epochs", 3)
+        learning_rate = self._get_config_value("learning_rate", 2e-5)
+        output_dir = self._get_config_value("output_dir", Path("output"))
+        logging_dir = self._get_config_value("logging_dir", None)
+
+        # load model
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name, num_labels=num_labels
+        )
+
+        # create lightning module
+        self.lightning_module = create_lightning_module(model, learning_rate)
+
+        # create callbacks
+        callbacks = [
+            pl.callbacks.ModelCheckpoint(
+                monitor="val_loss",
+                dirpath=output_dir,
+                filename="best-{epoch:02d}-{val_loss:.2f}",
+            ),
+            pl.callbacks.EarlyStopping(monitor="val_loss", patience=3),
+        ]
+
+        # create logger
+        logger = None
+        if logging_dir:
+            logger = pl.loggers.TensorBoardLogger(str(logging_dir))
+
+        # create trainer
+        trainer = pl.Trainer(
+            max_epochs=num_epochs,
+            accelerator="auto",
+            devices="auto",
+            logger=logger,
+            callbacks=callbacks,
+        )
+
+        # train
+        trainer.fit(
+            self.lightning_module,
+            train_dataloaders=train_data,
+            val_dataloaders=eval_data,
+        )
+
+        # evaluate
+        metrics: dict[str, float] = {}
+        if eval_data is not None:
+            eval_results = trainer.validate(
+                self.lightning_module, dataloaders=eval_data
+            )
+            if eval_results:
+                metrics = {k: float(v) for k, v in eval_results[0].items()}
+
+        training_time = time.time() - start_time
+
+        # get best checkpoint path
+        best_checkpoint = None
+        if hasattr(trainer.checkpoint_callback, "best_model_path"):
+            best_checkpoint_str = trainer.checkpoint_callback.best_model_path
+            if best_checkpoint_str:
+                best_checkpoint = Path(best_checkpoint_str)
+
+        # create metadata
+        config_dict = (
+            self.config
+            if isinstance(self.config, dict)
+            else (
+                self.config.model_dump() if hasattr(self.config, "model_dump") else {}
+            )
+        )
+
+        metadata = ModelMetadata(
+            model_name=model_name,
+            framework="pytorch_lightning",
+            training_config=config_dict,
+            training_data_path=Path("train.json"),
+            eval_data_path=Path("eval.json") if eval_data else None,
+            metrics=metrics,
+            best_checkpoint=best_checkpoint,
+            training_time=training_time,
+            training_timestamp=format_iso8601(now_iso8601()),
+        )
+
+        return metadata
+
+    def save_model(self, output_dir: Path, metadata: ModelMetadata) -> None:
+        """Save model and metadata to disk.
+
+        Saves the Lightning module state dict and training metadata as JSON.
+
+        Parameters
+        ----------
+        output_dir
+            Directory to save model checkpoint and metadata JSON file.
+        metadata
+            Training metadata to save alongside the model.
+
+        Examples
+        --------
+        >>> trainer = PyTorchLightningTrainer({})  # doctest: +SKIP
+        >>> trainer.save_model(Path("output"), metadata)  # doctest: +SKIP
+        """
+        import torch  # noqa: PLC0415
+
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # save lightning checkpoint
+        if self.lightning_module is not None:
+            torch.save(
+                self.lightning_module.state_dict(),
+                output_dir / "lightning_model.pt",
+            )
+
+        # save metadata
+        with open(output_dir / "metadata.json", "w") as f:
+            metadata_dict = metadata.model_dump()
+            json.dump(metadata_dict, f, indent=2, default=str)
+
+    def load_model(self, model_dir: Path) -> pl.LightningModule | None:
+        """Load a saved model from disk.
+
+        Parameters
+        ----------
+        model_dir
+            Directory containing the saved Lightning model state dict.
+
+        Returns
+        -------
+        pl.LightningModule | None
+            The Lightning module with loaded weights, or None if no module
+            has been initialized.
+
+        Examples
+        --------
+        >>> trainer = PyTorchLightningTrainer({})  # doctest: +SKIP
+        >>> model = trainer.load_model(Path("saved_model"))  # doctest: +SKIP
+        """
+        import torch  # noqa: PLC0415
+
+        if self.lightning_module is not None:
+            self.lightning_module.load_state_dict(
+                torch.load(model_dir / "lightning_model.pt")
+            )
+        return self.lightning_module