feat!(tokenizer): ensure output is consistent across al tokenizers

Author: meilame-tayebjee

torchTextClassifiers/dataset/dataset.py

@@ -81,8 +81,8 @@ def collate_fn(self, batch):
             categorical_tensors = None
 
         return {
-            "input_ids": tokenize_output["input_ids"],
-            "attention_mask": tokenize_output["attention_mask"],
+            "input_ids": tokenize_output.input_ids,
+            "attention_mask": tokenize_output.attention_mask,
             "categorical_vars": categorical_tensors,
             "labels": labels_tensor,
         }

torchTextClassifiers/tokenizers/__init__.py

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

torchTextClassifiers/tokenizers/base.py

@@ -1,5 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional, Union
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
 
 try:
     from tokenizers import Tokenizer
@@ -10,6 +14,55 @@
     HAS_HF = False
 
 
+@dataclass
+class TokenizerOutput:
+    input_ids: torch.Tensor  # shape: (batch_size, seq_len)
+    attention_mask: torch.Tensor  # shape: (batch_size, seq_len)
+    offset_mapping: Optional[torch.Tensor] = None  # shape: (batch_size, seq_len, 2)
+    word_ids: Optional[np.ndarray] = None  # shape: (batch_size, seq_len)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TokenizerOutput":
+        return cls(**data)
+
+    def __post_init__(self):
+        # --- Basic type checks ---
+        if not isinstance(self.input_ids, torch.Tensor):
+            raise TypeError(f"token_ids must be a torch.Tensor, got {type(self.input_ids)}")
+        if not isinstance(self.attention_mask, torch.Tensor):
+            raise TypeError(
+                f"attention_mask must be a torch.Tensor, got {type(self.attention_mask)}"
+            )
+        if self.offset_mapping is not None and not isinstance(self.offset_mapping, torch.Tensor):
+            raise TypeError(
+                f"offset_mapping must be a torch.Tensor or None, got {type(self.offset_mapping)}"
+            )
+        if self.word_ids is not None and not isinstance(self.word_ids, np.ndarray):
+            raise TypeError(f"word_ids must be a numpy.ndarray or None, got {type(self.word_ids)}")
+
+        # --- Shape consistency checks ---
+        if self.input_ids.shape != self.attention_mask.shape:
+            raise ValueError(
+                f"Shape mismatch: token_ids {self.token_ids.shape} and attention_mask {self.attention_mask.shape}"
+            )
+
+        if self.offset_mapping is not None:
+            expected_shape = (*self.input_ids.shape, 2)
+            if self.offset_mapping.shape != expected_shape:
+                raise ValueError(
+                    f"offset_mapping should have shape {expected_shape}, got {self.offset_mapping.shape}"
+                )
+
+        if self.word_ids is not None:
+            if self.word_ids.shape != self.input_ids.shape:
+                raise ValueError(
+                    f"word_ids should have shape {self.input_ids.shape}, got {self.word_ids.shape}"
+                )
+
+
 class BaseTokenizer(ABC):
     def __init__(
         self, vocab_size: int, output_vectorized: bool = False, output_dim: Optional[int] = None
@@ -32,7 +85,7 @@ def __init__(
                 )
 
     @abstractmethod
-    def tokenize(self, text: Union[str, List[str]]) -> list:
+    def tokenize(self, text: Union[str, List[str]]) -> TokenizerOutput:
         """Tokenizes the raw input text into a list of tokens."""
         pass
 
@@ -72,14 +125,23 @@ def tokenize(
         # Pad to longest sequence if no output_dim is specified
         padding = True if self.output_dim is None else "max_length"
 
-        return self.tokenizer(
+        tokenize_output = self.tokenizer(
             text,
             padding=padding,
             return_tensors="pt",
             max_length=self.output_dim,
             return_offsets_mapping=return_offsets_mapping,
         )  # method from PreTrainedTokenizerFast
 
+        encoded_text = tokenize_output["input_ids"]
+
+        return TokenizerOutput(
+            input_ids=encoded_text,
+            attention_mask=tokenize_output["attention_mask"],
+            offset_mapping=tokenize_output.get("offset_mapping", None),
+            word_ids=np.array([tokenize_output.word_ids(i) for i in range(len(encoded_text))]),
+        )
+
     @classmethod
     def load_from_pretrained(cls, tokenizer_name: str, output_dim: Optional[int] = None):
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

torchTextClassifiers/torchTextClassifiers.py

@@ -28,7 +28,7 @@
     ClassificationHead,
     TextEmbedder,
 )
-from torchTextClassifiers.tokenizers import BaseTokenizer
+from torchTextClassifiers.tokenizers import BaseTokenizer, TokenizerOutput
 
 logger = logging.getLogger(__name__)
 
@@ -474,8 +474,13 @@ def predict(
             text.tolist(), return_offsets_mapping=return_offsets_mapping
         )
 
-        encoded_text = tokenize_output["input_ids"]  # (batch_size, seq_len)
-        attention_mask = tokenize_output["attention_mask"]  # (batch_size, seq_len)
+        if not isinstance(tokenize_output, TokenizerOutput):
+            raise TypeError(
+                f"Expected TokenizerOutput, got {type(tokenize_output)} from tokenizer.tokenize method."
+            )
+
+        encoded_text = tokenize_output.input_ids  # (batch_size, seq_len)
+        attention_mask = tokenize_output.attention_mask  # (batch_size, seq_len)
 
         if categorical_variables is not None:
             categorical_vars = torch.tensor(
@@ -511,10 +516,8 @@ def predict(
                 "prediction": predictions,
                 "confidence": confidence,
                 "attributions": all_attributions,
-                "offset_mapping": tokenize_output.get("offset_mapping", None),
-                "word_ids": np.array(
-                    [tokenize_output.word_ids(i) for i in range(len(encoded_text))]
-                ),
+                "offset_mapping": tokenize_output.offset_mapping,
+                "word_ids": tokenize_output.word_ids,
             }
         else:
             return {