torchtextclassifiers 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

torchTextClassifiers/model/components/text_embedder.py

@@ -0,0 +1,220 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import nn
+
+from torchTextClassifiers.model.components.attention import AttentionConfig, Block, norm
+
+
+@dataclass
+class TextEmbedderConfig:
+    vocab_size: int
+    embedding_dim: int
+    padding_idx: int
+    attention_config: Optional[AttentionConfig] = None
+
+
+class TextEmbedder(nn.Module):
+    def __init__(self, text_embedder_config: TextEmbedderConfig):
+        super().__init__()
+
+        self.config = text_embedder_config
+
+        self.attention_config = text_embedder_config.attention_config
+        if self.attention_config is not None:
+            self.attention_config.n_embd = text_embedder_config.embedding_dim
+
+        self.vocab_size = text_embedder_config.vocab_size
+        self.embedding_dim = text_embedder_config.embedding_dim
+        self.padding_idx = text_embedder_config.padding_idx
+
+        self.embedding_layer = nn.Embedding(
+            embedding_dim=self.embedding_dim,
+            num_embeddings=self.vocab_size,
+            padding_idx=self.padding_idx,
+        )
+
+        if self.attention_config is not None:
+            self.transformer = nn.ModuleDict(
+                {
+                    "h": nn.ModuleList(
+                        [
+                            Block(self.attention_config, layer_idx)
+                            for layer_idx in range(self.attention_config.n_layers)
+                        ]
+                    ),
+                }
+            )
+
+            head_dim = self.attention_config.n_embd // self.attention_config.n_head
+
+            if head_dim * self.attention_config.n_head != self.attention_config.n_embd:
+                raise ValueError("embedding_dim must be divisible by n_head.")
+
+            if self.attention_config.positional_encoding:
+                if head_dim % 2 != 0:
+                    raise ValueError(
+                        "embedding_dim / n_head must be even for rotary positional embeddings."
+                    )
+
+                if self.attention_config.sequence_len is None:
+                    raise ValueError(
+                        "sequence_len must be specified in AttentionConfig when positional_encoding is True."
+                    )
+
+                self.rotary_seq_len = self.attention_config.sequence_len * 10
+                cos, sin = self._precompute_rotary_embeddings(
+                    seq_len=self.rotary_seq_len, head_dim=head_dim
+                )
+
+                self.register_buffer(
+                    "cos", cos, persistent=False
+                )  # persistent=False means it's not saved to the checkpoint
+                self.register_buffer("sin", sin, persistent=False)
+
+    def init_weights(self):
+        self.apply(self._init_weights)
+
+        # zero out c_proj weights in all blocks
+        if self.attention_config is not None:
+            for block in self.transformer.h:
+                torch.nn.init.zeros_(block.mlp.c_proj.weight)
+                torch.nn.init.zeros_(block.attn.c_proj.weight)
+            # init the rotary embeddings
+            head_dim = self.attention_config.n_embd // self.attention_config.n_head
+            cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+            self.cos, self.sin = cos, sin
+        # Cast the embeddings from fp32 to bf16: optim can tolerate it and it saves memory: both in the model and the activations
+        if self.embedding_layer.weight.device.type == "cuda":
+            self.embedding_layer.to(dtype=torch.bfloat16)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # https://arxiv.org/pdf/2310.17813
+            fan_out = module.weight.size(0)
+            fan_in = module.weight.size(1)
+            std = 1.0 / math.sqrt(fan_in) * min(1.0, math.sqrt(fan_out / fan_in))
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=1.0)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        """Converts input token IDs to their corresponding embeddings."""
+
+        encoded_text = input_ids  # clearer name
+        if encoded_text.dtype != torch.long:
+            encoded_text = encoded_text.to(torch.long)
+
+        batch_size, seq_len = encoded_text.shape
+        batch_size_check, seq_len_check = attention_mask.shape
+
+        if batch_size != batch_size_check or seq_len != seq_len_check:
+            raise ValueError(
+                f"Input IDs and attention mask must have the same batch size and sequence length. "
+                f"Got input_ids shape {encoded_text.shape} and attention_mask shape {attention_mask.shape}."
+            )
+
+        token_embeddings = self.embedding_layer(
+            encoded_text
+        )  # (batch_size, seq_len, embedding_dim)
+
+        token_embeddings = norm(token_embeddings)
+
+        if self.attention_config is not None:
+            if self.attention_config.positional_encoding:
+                cos_sin = self.cos[:, :seq_len], self.sin[:, :seq_len]
+            else:
+                cos_sin = None
+
+            for block in self.transformer.h:
+                token_embeddings = block(token_embeddings, cos_sin)
+
+            token_embeddings = norm(token_embeddings)
+
+        text_embedding = self._get_sentence_embedding(
+            token_embeddings=token_embeddings, attention_mask=attention_mask
+        )
+
+        return text_embedding
+
+    def _get_sentence_embedding(
+        self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Compute sentence embedding from embedded tokens - "remove" second dimension.
+
+        Args (output from dataset collate_fn):
+            token_embeddings (torch.Tensor[Long]), shape (batch_size, seq_len, embedding_dim): Tokenized + padded text
+            attention_mask (torch.Tensor[Long]), shape (batch_size, seq_len): Attention mask indicating non-pad tokens
+        Returns:
+            torch.Tensor: Sentence embeddings, shape (batch_size, embedding_dim)
+        """
+
+        # average over non-pad token embeddings
+        # attention mask has 1 for non-pad tokens and 0 for pad token positions
+
+        # mask pad-tokens
+
+        if self.attention_config is not None:
+            if self.attention_config.aggregation_method is not None:
+                if self.attention_config.aggregation_method == "first":
+                    return token_embeddings[:, 0, :]
+                elif self.attention_config.aggregation_method == "last":
+                    lengths = attention_mask.sum(dim=1).clamp(min=1)  # last non-pad token index + 1
+                    return token_embeddings[
+                        torch.arange(token_embeddings.size(0)),
+                        lengths - 1,
+                        :,
+                    ]
+                else:
+                    if self.attention_config.aggregation_method != "mean":
+                        raise ValueError(
+                            f"Unknown aggregation method: {self.attention_config.aggregation_method}. Supported methods are 'mean', 'first', 'last'."
+                        )
+
+        assert self.attention_config is None or self.attention_config.aggregation_method == "mean"
+
+        mask = attention_mask.unsqueeze(-1).float()  # (batch_size, seq_len, 1)
+        masked_embeddings = token_embeddings * mask  # (batch_size, seq_len, embedding_dim)
+
+        sentence_embedding = masked_embeddings.sum(dim=1) / mask.sum(dim=1).clamp(
+            min=1.0
+        )  # avoid division by zero
+
+        sentence_embedding = torch.nan_to_num(sentence_embedding, 0.0)
+
+        return sentence_embedding
+
+    def __call__(self, *args, **kwargs):
+        out = super().__call__(*args, **kwargs)
+        if out.dim() != 2:
+            raise ValueError(
+                f"Output of {self.__class__.__name__}.forward must be 2D "
+                f"(got shape {tuple(out.shape)})"
+            )
+        return out
+
+    def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None):
+        # autodetect the device from model embeddings
+        if device is None:
+            device = next(self.parameters()).device
+
+        # stride the channels
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        # stride the time steps
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        # calculate the rotation frequencies at each (time, channel) pair
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()  # keep them in bfloat16
+        cos, sin = (
+            cos[None, :, None, :],
+            sin[None, :, None, :],
+        )  # add batch and head dims for later broadcasting
+
+        return cos, sin

torchTextClassifiers/model/lightning.py

@@ -0,0 +1,166 @@
+import pytorch_lightning as pl
+import torch
+from torchmetrics import Accuracy
+
+from .model import TextClassificationModel
+
+# ============================================================================
+# PyTorch Lightning Module
+# ============================================================================
+
+
+class TextClassificationModule(pl.LightningModule):
+    """Pytorch Lightning Module for FastTextModel."""
+
+    def __init__(
+        self,
+        model: TextClassificationModel,
+        loss,
+        optimizer,
+        optimizer_params,
+        scheduler,
+        scheduler_params,
+        scheduler_interval="epoch",
+        **kwargs,
+    ):
+        """
+        Initialize FastTextModule.
+
+        Args:
+            model: Model.
+            loss: Loss
+            optimizer: Optimizer
+            optimizer_params: Optimizer parameters.
+            scheduler: Scheduler.
+            scheduler_params: Scheduler parameters.
+            scheduler_interval: Scheduler interval.
+        """
+        super().__init__()
+        self.save_hyperparameters(ignore=["model", "loss"])
+
+        self.model = model
+        self.loss = loss
+        self.accuracy_fn = Accuracy(task="multiclass", num_classes=self.model.num_classes)
+        self.optimizer = optimizer
+        self.optimizer_params = optimizer_params
+        self.scheduler = scheduler
+        self.scheduler_params = scheduler_params
+        self.scheduler_interval = scheduler_interval
+
+    def forward(self, batch) -> torch.Tensor:
+        """
+        Perform forward-pass.
+
+        Args:
+            batch (List[torch.LongTensor]): Batch to perform forward-pass on.
+
+        Returns (torch.Tensor): Prediction.
+        """
+        return self.model(
+            input_ids=batch["input_ids"],
+            attention_mask=batch["attention_mask"],
+            categorical_vars=batch.get("categorical_vars", None),
+        )
+
+    def training_step(self, batch, batch_idx: int) -> torch.Tensor:
+        """
+        Training step.
+
+        Args:
+            batch (List[torch.LongTensor]): Training batch.
+            batch_idx (int): Batch index.
+
+        Returns (torch.Tensor): Loss tensor.
+        """
+
+        targets = batch["labels"]
+
+        outputs = self.forward(batch)
+        loss = self.loss(outputs, targets)
+        self.log("train_loss", loss, on_epoch=True, on_step=True, prog_bar=True)
+        accuracy = self.accuracy_fn(outputs, targets)
+        self.log("train_accuracy", accuracy, on_epoch=True, on_step=False, prog_bar=True)
+
+        torch.cuda.empty_cache()
+
+        return loss
+
+    def validation_step(self, batch, batch_idx: int):
+        """
+        Validation step.
+
+        Args:
+            batch (List[torch.LongTensor]): Validation batch.
+            batch_idx (int): Batch index.
+
+        Returns (torch.Tensor): Loss tensor.
+        """
+        targets = batch["labels"]
+
+        outputs = self.forward(batch)
+        loss = self.loss(outputs, targets)
+        self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True, sync_dist=True)
+
+        accuracy = self.accuracy_fn(outputs, targets)
+        self.log("val_accuracy", accuracy, on_epoch=True, on_step=False, prog_bar=True)
+        return loss
+
+    def test_step(self, batch, batch_idx: int):
+        """
+        Test step.
+
+        Args:
+            batch (List[torch.LongTensor]): Test batch.
+            batch_idx (int): Batch index.
+
+        Returns (torch.Tensor): Loss tensor.
+        """
+        targets = batch["labels"]
+
+        outputs = self.forward(batch)
+        loss = self.loss(outputs, targets)
+
+        accuracy = self.accuracy_fn(outputs, targets)
+
+        return loss, accuracy
+
+    def predict_step(self, batch, batch_idx: int = 0, dataloader_idx: int = 0):
+        """
+        Prediction step.
+
+        Args:
+            batch (List[torch.LongTensor]): Prediction batch.
+            batch_idx (int): Batch index.
+            dataloader_idx (int): Dataloader index.
+
+        Returns (torch.Tensor): Predictions.
+        """
+        outputs = self.forward(batch)
+        return outputs
+
+    def configure_optimizers(self):
+        """
+        Configure optimizer for Pytorch lighting.
+
+        Returns: Optimizer and scheduler for pytorch lighting.
+        """
+        optimizer = self.optimizer(self.parameters(), **self.optimizer_params)
+
+        if self.scheduler is None:
+            return optimizer
+
+        # Only use scheduler if it's not ReduceLROnPlateau or if we can ensure val_loss is available
+        # For complex training setups, sometimes val_loss is not available every epoch
+        if hasattr(self.scheduler, "__name__") and "ReduceLROnPlateau" in self.scheduler.__name__:
+            # For ReduceLROnPlateau, use train_loss as it's always available
+            scheduler = self.scheduler(optimizer, **self.scheduler_params)
+            scheduler_config = {
+                "scheduler": scheduler,
+                "monitor": "train_loss",
+                "interval": self.scheduler_interval,
+            }
+            return [optimizer], [scheduler_config]
+        else:
+            # For other schedulers (StepLR, etc.), no monitoring needed
+            scheduler = self.scheduler(optimizer, **self.scheduler_params)
+            return [optimizer], [scheduler]

torchTextClassifiers/model/model.py

@@ -0,0 +1,151 @@
+"""FastText model components.
+
+This module contains the PyTorch model, Lightning module, and dataset classes
+for FastText classification. Consolidates what was previously in pytorch_model.py,
+lightning_module.py, and dataset.py.
+"""
+
+import logging
+from typing import Annotated, Optional
+
+import torch
+from torch import nn
+
+from torchTextClassifiers.model.components import (
+    CategoricalForwardType,
+    CategoricalVariableNet,
+    ClassificationHead,
+    TextEmbedder,
+)
+
+logger = logging.getLogger(__name__)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    handlers=[logging.StreamHandler()],
+)
+
+
+# ============================================================================
+# PyTorch Model
+
+# It takes PyTorch tensors as input (not raw text!),
+# and it outputs raw not-softmaxed logits, not predictions
+# ============================================================================
+
+
+class TextClassificationModel(nn.Module):
+    """FastText Pytorch Model."""
+
+    def __init__(
+        self,
+        classification_head: ClassificationHead,
+        text_embedder: Optional[TextEmbedder] = None,
+        categorical_variable_net: Optional[CategoricalVariableNet] = None,
+    ):
+        """
+        Constructor for the FastTextModel class.
+
+        Args:
+            classification_head (ClassificationHead): The classification head module.
+            text_embedder (Optional[TextEmbedder]): The text embedding module.
+                If not provided, assumes that input text is already embedded (as tensors) and directly passed to the classification head.
+            categorical_variable_net (Optional[CategoricalVariableNet]): The categorical variable network module.
+                If not provided, assumes no categorical variables are used.
+        """
+        super().__init__()
+
+        self.text_embedder = text_embedder
+
+        self.categorical_variable_net = categorical_variable_net
+        if not self.categorical_variable_net:
+            logger.info("🔹 No categorical variable network provided; using only text embeddings.")
+
+        self.classification_head = classification_head
+
+        self._validate_component_connections()
+
+        self.num_classes = self.classification_head.num_classes
+
+        torch.nn.init.zeros_(self.classification_head.net.weight)
+        if self.text_embedder is not None:
+            self.text_embedder.init_weights()
+
+    def _validate_component_connections(self):
+        def _check_text_categorical_connection(self, text_embedder, cat_var_net):
+            if cat_var_net.forward_type == CategoricalForwardType.SUM_TO_TEXT:
+                if text_embedder.embedding_dim != cat_var_net.output_dim:
+                    raise ValueError(
+                        "Text embedding dimension must match categorical variable embedding dimension."
+                    )
+                self.expected_classification_head_input_dim = text_embedder.embedding_dim
+            else:
+                self.expected_classification_head_input_dim = (
+                    text_embedder.embedding_dim + cat_var_net.output_dim
+                )
+
+        if self.text_embedder:
+            if self.categorical_variable_net:
+                _check_text_categorical_connection(
+                    self, self.text_embedder, self.categorical_variable_net
+                )
+            else:
+                self.expected_classification_head_input_dim = self.text_embedder.embedding_dim
+
+            if self.expected_classification_head_input_dim != self.classification_head.input_dim:
+                raise ValueError(
+                    "Classification head input dimension does not match expected dimension from text embedder and categorical variable net."
+                )
+        else:
+            logger.warning(
+                "⚠️ No text embedder provided; assuming input text is already embedded or vectorized. Take care that the classification head input dimension matches the input text dimension."
+            )
+
+    def forward(
+        self,
+        input_ids: Annotated[torch.Tensor, "batch seq_len"],
+        attention_mask: Annotated[torch.Tensor, "batch seq_len"],
+        categorical_vars: Annotated[torch.Tensor, "batch num_cats"],
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Memory-efficient forward pass implementation.
+
+        Args: output from dataset collate_fn
+            input_ids (torch.Tensor[Long]), shape (batch_size, seq_len): Tokenized + padded text
+            attention_mask (torch.Tensor[int]), shape (batch_size, seq_len): Attention mask indicating non-pad tokens
+            categorical_vars (torch.Tensor[Long]): Additional categorical features, (batch_size, num_categorical_features)
+
+        Returns:
+            torch.Tensor: Model output scores for each class - shape (batch_size, num_classes)
+                Raw, not softmaxed.
+        """
+        encoded_text = input_ids  # clearer name
+        if self.text_embedder is None:
+            x_text = encoded_text.float()
+        else:
+            x_text = self.text_embedder(input_ids=encoded_text, attention_mask=attention_mask)
+
+        if self.categorical_variable_net:
+            x_cat = self.categorical_variable_net(categorical_vars)
+
+            if (
+                self.categorical_variable_net.forward_type
+                == CategoricalForwardType.AVERAGE_AND_CONCAT
+                or self.categorical_variable_net.forward_type
+                == CategoricalForwardType.CONCATENATE_ALL
+            ):
+                x_combined = torch.cat((x_text, x_cat), dim=1)
+            else:
+                assert (
+                    self.categorical_variable_net.forward_type == CategoricalForwardType.SUM_TO_TEXT
+                )
+                x_combined = x_text + x_cat
+        else:
+            x_combined = x_text
+
+        logits = self.classification_head(x_combined)
+
+        return logits

torchTextClassifiers/tokenizers/WordPiece.py

@@ -0,0 +1,92 @@
+import logging
+import os
+from typing import List, Optional
+
+from torchTextClassifiers.tokenizers import HAS_HF, HuggingFaceTokenizer
+
+if not HAS_HF:
+    raise ImportError(
+        "The HuggingFace dependencies are needed to use this tokenizer. Please run 'uv add torchTextClassifiers --extra huggingface."
+    )
+else:
+    from tokenizers import (
+        Tokenizer,
+        decoders,
+        models,
+        normalizers,
+        pre_tokenizers,
+        processors,
+        trainers,
+    )
+    from transformers import PreTrainedTokenizerFast
+
+logger = logging.getLogger(__name__)
+
+
+class WordPieceTokenizer(HuggingFaceTokenizer):
+    def __init__(self, vocab_size: int, trained: bool = False, output_dim: Optional[int] = None):
+        """Largely inspired by https://huggingface.co/learn/llm-course/chapter6/8"""
+
+        super().__init__(vocab_size=vocab_size, output_dim=output_dim)
+
+        self.unk_token = "[UNK]"
+        self.pad_token = "[PAD]"
+        self.cls_token = "[CLS]"
+        self.sep_token = "[SEP]"
+        self.special_tokens = [
+            self.unk_token,
+            self.pad_token,
+            self.cls_token,
+            self.sep_token,
+        ]
+        self.vocab_size = vocab_size
+        self.context_size = output_dim
+
+        self.tokenizer = Tokenizer(models.WordPiece(unk_token=self.unk_token))
+
+        self.tokenizer.normalizer = normalizers.BertNormalizer(
+            lowercase=True
+        )  # NFD, lowercase, strip accents - BERT style
+
+        self.tokenizer.pre_tokenizer = (
+            pre_tokenizers.BertPreTokenizer()
+        )  # split on whitespace and punctuation - BERT style
+        self.trained = trained
+
+    def _post_training(self):
+        if not self.trained:
+            raise RuntimeError(
+                "Tokenizer must be trained before applying post-training configurations."
+            )
+
+        self.tokenizer.post_processor = processors.BertProcessing(
+            (self.cls_token, self.tokenizer.token_to_id(self.cls_token)),
+            (self.sep_token, self.tokenizer.token_to_id(self.sep_token)),
+        )
+        self.tokenizer.decoder = decoders.WordPiece(prefix="##")
+        self.padding_idx = self.tokenizer.token_to_id("[PAD]")
+        self.tokenizer.enable_padding(pad_id=self.padding_idx, pad_token="[PAD]")
+
+        self.tokenizer = PreTrainedTokenizerFast(tokenizer_object=self.tokenizer)
+        self.vocab_size = len(self.tokenizer)
+
+    def train(
+        self, training_corpus: List[str], save_path: str = None, filesystem=None, s3_save_path=None
+    ):
+        trainer = trainers.WordPieceTrainer(
+            vocab_size=self.vocab_size,
+            special_tokens=self.special_tokens,
+        )
+        self.tokenizer.train_from_iterator(training_corpus, trainer=trainer)
+        self.trained = True
+        self._post_training()
+
+        if save_path:
+            self.tokenizer.save(save_path)
+            logger.info(f"💾 Tokenizer saved at {save_path}")
+            if filesystem and s3_save_path:
+                parent_dir = os.path.dirname(save_path)
+                if not filesystem.exists(parent_dir):
+                    filesystem.mkdirs(parent_dir)
+                filesystem.put(save_path, s3_save_path)
+                logger.info(f"💾 Tokenizer uploaded to S3 at {s3_save_path}")

torchTextClassifiers/tokenizers/__init__.py

@@ -0,0 +1,10 @@
+from .base import (
+    HAS_HF as HAS_HF,
+)
+from .base import BaseTokenizer as BaseTokenizer
+from .base import (
+    HuggingFaceTokenizer as HuggingFaceTokenizer,
+)
+from .base import TokenizerOutput as TokenizerOutput
+from .ngram import NGramTokenizer as NGramTokenizer
+from .WordPiece import WordPieceTokenizer as WordPieceTokenizer