torchtextclassifiers 1.0.0__tar.gz → 1.0.2__tar.gz

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: torchtextclassifiers
-Version: 1.0.0
+Version: 1.0.2
 Summary: A text classification toolkit to easily build, train and evaluate deep learning text classifiers using PyTorch.
 Keywords: fastText,text classification,NLP,automatic coding,deep learning
 Author: Cédric Couralet, Meilame Tayebjee

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "pytorch-lightning>=2.4.0",
 ]
 requires-python = ">=3.11"
-version="1.0.0"
+version="1.0.2"
 
 
 [dependency-groups]

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/torchTextClassifiers/model/components/text_embedder.py

@@ -23,6 +23,9 @@ class TextEmbedder(nn.Module):
         self.config = text_embedder_config
 
         self.attention_config = text_embedder_config.attention_config
+        if isinstance(self.attention_config, dict):
+            self.attention_config = AttentionConfig(**self.attention_config)
+
         if self.attention_config is not None:
             self.attention_config.n_embd = text_embedder_config.embedding_dim
 

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/torchTextClassifiers/model/lightning.py

@@ -36,7 +36,7 @@ class TextClassificationModule(pl.LightningModule):
             scheduler_interval: Scheduler interval.
         """
         super().__init__()
-        self.save_hyperparameters(ignore=["model", "loss"])
+        self.save_hyperparameters(ignore=["model"])
 
         self.model = model
         self.loss = loss

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/torchTextClassifiers/tokenizers/__init__.py

@@ -7,4 +7,6 @@ from .base import (
 )
 from .base import TokenizerOutput as TokenizerOutput
 from .ngram import NGramTokenizer as NGramTokenizer
-from .WordPiece import WordPieceTokenizer as WordPieceTokenizer
+
+if HAS_HF:
+    from .WordPiece import WordPieceTokenizer as WordPieceTokenizer

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/torchTextClassifiers/torchTextClassifiers.py

@@ -1,6 +1,8 @@
 import logging
+import pickle
 import time
 from dataclasses import asdict, dataclass, field
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 try:
@@ -75,6 +77,7 @@ class TrainingConfig:
     trainer_params: Optional[dict] = None
     optimizer_params: Optional[dict] = None
     scheduler_params: Optional[dict] = None
+    save_path: Optional[str] = "my_ttc"
 
     def to_dict(self) -> Dict[str, Any]:
         data = asdict(self)
@@ -198,6 +201,12 @@ class torchTextClassifiers:
         - Model training with early stopping
         - Best model loading after training
 
+        Note on Checkpoints:
+            After training, the best model checkpoint is automatically loaded.
+            This checkpoint contains the full training state (model weights,
+            optimizer, and scheduler state). Loading uses weights_only=False
+            as the checkpoint is self-generated and trusted.
+
         Args:
             X_train: Training input data
             y_train: Training labels
@@ -356,15 +365,20 @@ class torchTextClassifiers:
             logger.info(f"Training completed in {end - start:.2f} seconds.")
 
         best_model_path = trainer.checkpoint_callback.best_model_path
+        self.checkpoint_path = best_model_path
 
         self.lightning_module = TextClassificationModule.load_from_checkpoint(
             best_model_path,
             model=self.pytorch_model,
             loss=training_config.loss,
+            weights_only=False,  # Required: checkpoint contains optimizer/scheduler state
         )
 
         self.pytorch_model = self.lightning_module.model.to(self.device)
 
+        self.save_path = training_config.save_path
+        self.save(self.save_path)
+
         self.lightning_module.eval()
 
     def _check_XY(self, X: np.ndarray, Y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
@@ -569,6 +583,122 @@ class torchTextClassifiers:
                 "confidence": confidence,
             }
 
+    def save(self, path: Union[str, Path]) -> None:
+        """Save the complete torchTextClassifiers instance to disk.
+
+        This saves:
+        - Model configuration
+        - Tokenizer state
+        - PyTorch Lightning checkpoint (if trained)
+        - All other instance attributes
+
+        Args:
+            path: Directory path where the model will be saved
+
+        Example:
+            >>> ttc = torchTextClassifiers(tokenizer, model_config)
+            >>> ttc.train(X_train, y_train, training_config)
+            >>> ttc.save("my_model")
+        """
+        path = Path(path)
+        path.mkdir(parents=True, exist_ok=True)
+
+        # Save the checkpoint if model has been trained
+        checkpoint_path = None
+        if hasattr(self, "lightning_module"):
+            checkpoint_path = path / "model_checkpoint.ckpt"
+            # Save the current state as a checkpoint
+            trainer = pl.Trainer()
+            trainer.strategy.connect(self.lightning_module)
+            trainer.save_checkpoint(checkpoint_path)
+
+        # Prepare metadata to save
+        metadata = {
+            "model_config": self.model_config.to_dict(),
+            "ragged_multilabel": self.ragged_multilabel,
+            "vocab_size": self.vocab_size,
+            "embedding_dim": self.embedding_dim,
+            "categorical_vocabulary_sizes": self.categorical_vocabulary_sizes,
+            "num_classes": self.num_classes,
+            "checkpoint_path": str(checkpoint_path) if checkpoint_path else None,
+            "device": str(self.device) if hasattr(self, "device") else None,
+        }
+
+        # Save metadata
+        with open(path / "metadata.pkl", "wb") as f:
+            pickle.dump(metadata, f)
+
+        # Save tokenizer
+        tokenizer_path = path / "tokenizer.pkl"
+        with open(tokenizer_path, "wb") as f:
+            pickle.dump(self.tokenizer, f)
+
+        logger.info(f"Model saved successfully to {path}")
+
+    @classmethod
+    def load(cls, path: Union[str, Path], device: str = "auto") -> "torchTextClassifiers":
+        """Load a torchTextClassifiers instance from disk.
+
+        Args:
+            path: Directory path where the model was saved
+            device: Device to load the model on ('auto', 'cpu', 'cuda', etc.)
+
+        Returns:
+            Loaded torchTextClassifiers instance
+
+        Example:
+            >>> loaded_ttc = torchTextClassifiers.load("my_model")
+            >>> predictions = loaded_ttc.predict(X_test)
+        """
+        path = Path(path)
+
+        if not path.exists():
+            raise FileNotFoundError(f"Model directory not found: {path}")
+
+        # Load metadata
+        with open(path / "metadata.pkl", "rb") as f:
+            metadata = pickle.load(f)
+
+        # Load tokenizer
+        with open(path / "tokenizer.pkl", "rb") as f:
+            tokenizer = pickle.load(f)
+
+        # Reconstruct model_config
+        model_config = ModelConfig.from_dict(metadata["model_config"])
+
+        # Create instance
+        instance = cls(
+            tokenizer=tokenizer,
+            model_config=model_config,
+            ragged_multilabel=metadata["ragged_multilabel"],
+        )
+
+        # Set device
+        if device == "auto":
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            device = torch.device(device)
+        instance.device = device
+
+        # Load checkpoint if it exists
+        if metadata["checkpoint_path"]:
+            checkpoint_path = path / "model_checkpoint.ckpt"
+            if checkpoint_path.exists():
+                # Load the checkpoint with weights_only=False since it's our own trusted checkpoint
+                instance.lightning_module = TextClassificationModule.load_from_checkpoint(
+                    str(checkpoint_path),
+                    model=instance.pytorch_model,
+                    weights_only=False,
+                )
+                instance.pytorch_model = instance.lightning_module.model.to(device)
+                instance.checkpoint_path = str(checkpoint_path)
+                logger.info(f"Model checkpoint loaded from {checkpoint_path}")
+            else:
+                logger.warning(f"Checkpoint file not found at {checkpoint_path}")
+
+        logger.info(f"Model loaded successfully from {path}")
+        return instance
+
     def __repr__(self):
         model_type = (
             self.lightning_module.__repr__()

{torchtextclassifiers-1.0.0 → torchtextclassifiers-1.0.2}/torchTextClassifiers/utilities/plot_explainability.py

@@ -53,8 +53,18 @@ def map_attributions_to_char(attributions, offsets, text):
         np.exp(attributions_per_char), axis=1, keepdims=True
     )  # softmax normalization
 
+def get_id_to_word(text, word_ids, offsets):
+    words = {}
+    for idx, word_id in enumerate(word_ids):
+        if word_id is None:
+            continue
+        start, end = offsets[idx]
+        words[int(word_id)] = text[start:end]
+    
+    return words
+
 
-def map_attributions_to_word(attributions, word_ids):
+def map_attributions_to_word(attributions, text, word_ids, offsets):
     """
     Maps token-level attributions to word-level attributions based on word IDs.
     Args:
@@ -69,8 +79,9 @@ def map_attributions_to_word(attributions, word_ids):
         np.ndarray: Array of shape (top_k, num_words) containing word-level attributions.
             num_words is the number of unique words in the original text.
     """
-
+    
     word_ids = np.array(word_ids)
+    words = get_id_to_word(text, word_ids, offsets)
 
     # Convert None to -1 for easier processing (PAD tokens)
     word_ids_int = np.array([x if x is not None else -1 for x in word_ids], dtype=int)
@@ -99,7 +110,7 @@ def map_attributions_to_word(attributions, word_ids):
         )  # zero-out non-matching tokens and sum attributions for all tokens belonging to the same word
 
     # assert word_attributions.sum(axis=1) == attributions.sum(axis=1), "Sum of word attributions per top_k must equal sum of token attributions per top_k."
-    return np.exp(word_attributions) / np.sum(
+    return words, np.exp(word_attributions) / np.sum(
         np.exp(word_attributions), axis=1, keepdims=True
     )  # softmax normalization
 
@@ -131,7 +142,7 @@ def plot_attributions_at_char(
         fig, ax = plt.subplots(figsize=figsize)
         ax.bar(range(len(text)), attributions_per_char[i])
         ax.set_xticks(np.arange(len(text)))
-        ax.set_xticklabels(list(text), rotation=90)
+        ax.set_xticklabels(list(text), rotation=45)
         title = titles[i] if titles is not None else f"Attributions for Top {i+1} Prediction"
         ax.set_title(title)
         ax.set_xlabel("Characters in Text")
@@ -142,7 +153,7 @@ def plot_attributions_at_char(
 
 
 def plot_attributions_at_word(
-    text, attributions_per_word, figsize=(10, 2), titles: Optional[List[str]] = None
+    text, words, attributions_per_word, figsize=(10, 2), titles: Optional[List[str]] = None
 ):
     """
     Plots word-level attributions as a heatmap.
@@ -159,14 +170,13 @@ def plot_attributions_at_word(
             "matplotlib is required for plotting. Please install it to use this function."
         )
 
-    words = text.split()
     top_k = attributions_per_word.shape[0]
     all_plots = []
     for i in range(top_k):
         fig, ax = plt.subplots(figsize=figsize)
         ax.bar(range(len(words)), attributions_per_word[i])
         ax.set_xticks(np.arange(len(words)))
-        ax.set_xticklabels(words, rotation=90)
+        ax.set_xticklabels(words, rotation=45)
         title = titles[i] if titles is not None else f"Attributions for Top {i+1} Prediction"
         ax.set_title(title)
         ax.set_xlabel("Words in Text")