libmultilabel 0.6.2__tar.gz → 0.7.0__tar.gz

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: libmultilabel
-Version: 0.6.2
-Summary: A library for multi-label text classification
+Version: 0.7.0
+Summary: A library for multi-class and multi-label classification
 Home-page: https://github.com/ASUS-AICS/LibMultiLabel
 Author: LibMultiLabel Team
 License: MIT License

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/README.md

@@ -1,4 +1,4 @@
-# LibMultiLabel — a Library for Multi-class and Multi-label Text Classification
+# LibMultiLabel — a Library for Multi-class and Multi-label Classification
 
 LibMultiLabel is a library for binary, multi-class, and multi-label classification. It has the following functionalities
 

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/linear/linear.py

@@ -17,10 +17,13 @@ __all__ = [
     "predict_values",
     "get_topk_labels",
     "get_positive_labels",
+    "FlatModel",
 ]
 
 
 class FlatModel:
+    """A model returned from a training function."""
+
     def __init__(
         self,
         name: str,
@@ -619,7 +622,7 @@ def train_binary_and_multiclass(
 
 
 def predict_values(model, x: sparse.csr_matrix) -> np.ndarray:
-    """Calculates the decision values associated with x.
+    """Calculates the decision values associated with x, equivalent to model.predict_values(x).
 
     Args:
         model: A model returned from a training function.

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/linear/tree.py

@@ -10,7 +10,7 @@ from tqdm import tqdm
 
 from . import linear
 
-__all__ = ["train_tree"]
+__all__ = ["train_tree", "TreeModel"]
 
 
 class Node:
@@ -38,6 +38,8 @@ class Node:
 
 
 class TreeModel:
+    """A model returned from train_tree."""
+
     def __init__(
         self,
         root: Node,

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/linear/utils.py

@@ -48,6 +48,7 @@ def save_pipeline(checkpoint_dir: str, preprocessor: Preprocessor, model):
                 "model": model,
             },
             f,
+            protocol=pickle.HIGHEST_PROTOCOL,
         )
 
 

libmultilabel-0.7.0/libmultilabel/nn/attentionxml.py

@@ -0,0 +1,800 @@
+from __future__ import annotations
+
+import logging
+from functools import partial
+from pathlib import Path
+from typing import Generator, Sequence, Optional
+
+import numpy as np
+import torch
+from lightning import Trainer
+from numpy import ndarray
+from scipy.sparse import csr_matrix, csc_matrix, issparse
+from scipy.special import expit
+from sklearn.preprocessing import MultiLabelBinarizer, normalize
+from torch import Tensor, is_tensor
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+from lightning.pytorch.callbacks import ModelCheckpoint
+
+from .data_utils import UNK
+from ..common_utils import dump_log
+from ..linear.preprocessor import Preprocessor
+from ..nn import networks
+from ..nn.model import Model
+
+__all__ = ["PLTTrainer"]
+
+from .nn_utils import init_trainer, init_model
+
+logger = logging.getLogger(__name__)
+
+
+class PLTTrainer:
+    CHECKPOINT_NAME = "model_"
+
+    def __init__(
+        self,
+        config,
+        classes: Optional[list] = None,
+        embed_vecs: Optional[Tensor] = None,
+        word_dict: Optional[dict] = None,
+    ):
+        # The number of levels is set to 2. In other words, there will be 2 models
+        self.multiclass = config.multiclass
+        if self.multiclass:
+            raise ValueError(
+                "The label space of multi-class datasets is usually not large, so PLT training is unnecessary."
+                "Please consider other methods."
+                "If you have a multi-class set with numerous labels, please let us know"
+            )
+
+        # cluster
+        self.cluster_size = config.cluster_size
+        # predict the top k clusters for deciding relevant/irrelevant labels of each instance in level 1 model training
+        self.beam_width = config.beam_width
+
+        # dataset meta info
+        self.embed_vecs = embed_vecs
+        self.word_dict = word_dict
+        self.classes = classes
+        self.max_seq_length = config.max_seq_length
+        self.num_classes = len(classes)
+
+        # multilabel binarizer fitted to the datasets
+        self.binarizer = None
+
+        # cluster meta info
+        self.cluster_size = config.cluster_size
+
+        # network parameters
+        self.network_config = config.network_config
+        self.init_weight = "xavier_uniform"  # AttentionXML-specific setting
+        self.loss_function = config.loss_function
+
+        # optimizer parameters
+        self.optimizer = config.optimizer
+        self.learning_rate = config.learning_rate
+        self.momentum = config.momentum
+        self.weight_decay = config.weight_decay
+        # learning rate scheduler
+        self.lr_scheduler = config.lr_scheduler
+        self.scheduler_config = config.scheduler_config
+
+        # Trainer parameters
+        self.use_cpu = config.cpu
+        self.accelerator = "cpu" if self.use_cpu else "gpu"
+        self.devices = 1
+        self.num_nodes = 1
+        self.epochs = config.epochs
+        self.limit_train_batches = config.limit_train_batches
+        self.limit_val_batches = config.limit_val_batches
+        self.limit_test_batches = config.limit_test_batches
+
+        # callbacks
+        self.silent = config.silent
+        # EarlyStopping
+        self.early_stopping_metric = config.early_stopping_metric
+        self.patience = config.patience
+        # ModelCheckpoint
+        self.val_metric = config.val_metric
+        self.checkpoint_dir = Path(config.checkpoint_dir)
+
+        self.metrics = config.monitor_metrics
+        self.metric_threshold = config.metric_threshold
+        self.monitor_metrics = config.monitor_metrics
+
+        # dataloader parameters
+        # whether shuffle the training dataset or not during the training process
+        self.shuffle = config.shuffle
+        pin_memory = True if self.accelerator == "gpu" else False
+        # training DataLoader
+        self.dataloader = partial(
+            DataLoader,
+            batch_size=config.batch_size,
+            num_workers=config.data_workers,
+            pin_memory=pin_memory,
+        )
+        # evaluation DataLoader
+        self.eval_dataloader = partial(
+            self.dataloader,
+            batch_size=config.eval_batch_size,
+        )
+
+        # predict
+        self.save_k_predictions = config.save_k_predictions
+
+        # save path
+        self.log_path = config.log_path
+        self.predict_out_path = config.predict_out_path
+        self.config = config
+
+    def label2cluster(self, cluster_mapping, *labels) -> Generator[csr_matrix, ...]:
+        """Map labels to their corresponding clusters in CSR sparse format.
+        Notice that this function deals with SPARSE matrix.
+        Assume there are 6 labels clustered as [(0, 1), (2, 3), (4, 5)]. Here (0, 1) is cluster with index 0 and so on.
+        Given the ground-truth labels, [0, 1, 4], the resulting clusters are [0, 2].
+
+        Args:
+            cluster_mapping (np.ndarray): mapping from clusters generated by build_label_tree to labels .
+            *labels (csr_matrix): labels in CSR sparse format.
+
+        Returns:
+            Generator[csr_matrix]: resulting clusters converted from labels in CSR sparse format
+        """
+        mapping = np.empty(self.num_classes, dtype=np.uint32)
+        for idx, clusters in enumerate(cluster_mapping):
+            mapping[clusters] = idx
+
+        def _label2cluster(label: csr_matrix) -> csr_matrix:
+            row = []
+            col = []
+            data = []
+            for i in range(label.shape[0]):
+                # n include all mapped ancestor clusters
+                n = np.unique(mapping[label.indices[label.indptr[i] : label.indptr[i + 1]]])
+                row += [i] * len(n)
+                col += n.tolist()
+                data += [1] * len(n)
+            return csr_matrix((data, (row, col)), shape=(label.shape[0], len(cluster_mapping)))
+
+        return (_label2cluster(label) for label in labels)
+
+    @staticmethod
+    def cluster2label(cluster_mapping, clusters, cluster_scores=None):
+        """Expand clusters to their corresponding labels and, if available, assign scores to each label.
+        Labels inside the same cluster have the same scores. This function is applied to predictions from model 0.
+        Notice that the behaviors of this function are different from label2cluster.
+        Also notice that this function deals with DENSE matrix.
+
+        Args:
+            cluster_mapping (np.ndarray): mapping from clusters generated by build_label_tree to labels .
+            clusters (np.ndarray): predicted clusters from model 0.
+            cluster_scores (Optional: np.ndarray): predicted scores of each cluster from model 0.
+
+        Returns:
+            Generator[np.ndarray]: resulting labels expanded from clusters
+        """
+
+        labels_selected = []
+
+        if cluster_scores is not None:
+            # label_scores are corresponding scores for selected labels and
+            # shape: (len(x), cluster_size * top_k)
+            label_scores = []
+            for score, cluster in zip(cluster_scores, clusters):
+                label_scores += [np.repeat(score, [len(labels) for labels in cluster_mapping[cluster]])]
+                labels_selected += [np.concatenate(cluster_mapping[cluster])]
+            return labels_selected, label_scores
+        else:
+            labels_selected = [np.concatenate(cluster_mapping[cluster]) for cluster in clusters]
+            return labels_selected
+
+    def fit(self, datasets):
+        """fit model to the training dataset
+
+        Args:
+            datasets: dict containing training, validation, and/or test datasets
+        """
+        if self.get_best_model_path(level=1).exists():
+            return
+
+        # AttentionXML-specific data preprocessing
+        train_val_dataset = datasets["train"] + datasets["val"]
+        train_val_dataset = {
+            "x": [" ".join(i["text"]) for i in train_val_dataset],
+            "y": [i["label"] for i in train_val_dataset],
+        }
+
+        # Preprocessor does tf-idf vectorization and multilabel binarization
+        # For details, see libmultilabel.linear.preprocessor.Preprocessor
+        preprocessor = Preprocessor()
+        datasets_temp = {"data_format": "txt", "train": train_val_dataset, "classes": self.classes}
+        # Preprocessor requires the input dictionary to has a key named "train" and will return a new dictionary with
+        # the same key.
+        train_val_dataset_tf = preprocessor.fit_transform(datasets_temp)["train"]
+        # save binarizer for testing
+        self.binarizer = preprocessor.binarizer
+
+        train_x = self.reformat_text(datasets["train"])
+        val_x = self.reformat_text(datasets["val"])
+
+        train_y = train_val_dataset_tf["y"][: len(datasets["train"])]
+        val_y = train_val_dataset_tf["y"][len(datasets["train"]) :]
+
+        # clusters are saved to the disk so that users doesn't need to provide the original training data when they want
+        # to do predicting solely
+        build_label_tree(
+            sparse_x=train_val_dataset_tf["x"],
+            sparse_y=train_val_dataset_tf["y"],
+            cluster_size=self.cluster_size,
+            output_dir=self.checkpoint_dir,
+        )
+
+        clusters = np.load(self.get_cluster_path(), allow_pickle=True)
+
+        # map each y to the parent cluster indices
+        train_y_clustered, val_y_clustered = self.label2cluster(clusters, train_y, val_y)
+
+        trainer = init_trainer(
+            self.checkpoint_dir,
+            epochs=self.epochs,
+            patience=self.patience,
+            early_stopping_metric=self.early_stopping_metric,
+            val_metric=self.val_metric,
+            silent=self.silent,
+            use_cpu=self.use_cpu,
+            limit_train_batches=self.limit_train_batches,
+            limit_val_batches=self.limit_val_batches,
+            limit_test_batches=self.limit_test_batches,
+            save_checkpoints=True,
+        )
+        trainer.checkpoint_callback.filename = f"{self.CHECKPOINT_NAME}0"
+
+        train_dataloader = self.dataloader(PlainDataset(train_x, train_y_clustered), shuffle=self.shuffle)
+        val_dataloader = self.dataloader(PlainDataset(val_x, val_y_clustered))
+
+        best_model_path = self.get_best_model_path(level=0)
+        if not best_model_path.exists():
+            model_0 = init_model(
+                model_name="AttentionXML_0",
+                network_config=self.network_config,
+                classes=clusters,
+                word_dict=self.word_dict,
+                embed_vecs=self.embed_vecs,
+                init_weight=self.init_weight,
+                log_path=self.log_path,
+                learning_rate=self.learning_rate,
+                optimizer=self.optimizer,
+                momentum=self.momentum,
+                weight_decay=self.weight_decay,
+                lr_scheduler=self.lr_scheduler,
+                scheduler_config=self.scheduler_config,
+                val_metric=self.val_metric,
+                metric_threshold=self.metric_threshold,
+                monitor_metrics=self.monitor_metrics,
+                multiclass=self.multiclass,
+                loss_function=self.loss_function,
+                silent=self.silent,
+                save_k_predictions=self.beam_width,
+            )
+
+            logger.info(f"Training level 0. Number of clusters: {len(clusters)}")
+            trainer.fit(model_0, train_dataloader, val_dataloader)
+            logger.info(f"Finish training level 0")
+
+        logger.info(f"Best model loaded from {best_model_path}")
+        model_0 = Model.load_from_checkpoint(best_model_path)
+
+        logger.info(
+            f"Predicting clusters by level-0 model. We then select {self.beam_width} clusters for each instance and "
+            f"extract labels from these clusters for level 1 training."
+        )
+        # load training and validation data and predict corresponding level 0 clusters
+        train_dataloader = self.dataloader(PlainDataset(train_x))
+        val_dataloader = self.dataloader(PlainDataset(val_x))
+
+        train_pred = trainer.predict(model_0, train_dataloader)
+        val_pred = trainer.predict(model_0, val_dataloader)
+
+        train_clusters_pred = np.vstack([i["top_k_pred"] for i in train_pred])
+        val_scores_pred = expit(np.vstack([i["top_k_pred_scores"] for i in val_pred]))
+        val_clusters_pred = np.vstack([i["top_k_pred"] for i in val_pred])
+
+        train_clusters_selected = np.empty((len(train_x), self.beam_width), dtype=np.uint)
+        for i, ys in enumerate(tqdm(train_clusters_pred, leave=False, desc="Sampling clusters")):
+            # relevant clusters are positive
+            pos = set(train_y_clustered.indices[train_y_clustered.indptr[i] : train_y_clustered.indptr[i + 1]])
+            # Select relevant clusters first. Then from top-predicted clusters, sequentially include them until
+            # cluster number reaches beam width
+            if len(pos) <= self.beam_width:
+                selected = pos
+                for y in ys:
+                    y = y.item()
+                    if len(selected) == self.beam_width:
+                        break
+                    selected.add(y)
+            # Regard positive (true) label as samples iff they appear in the predicted labels
+            # if the number of positive labels is more than top_k. If samples are not of length top_k
+            # add unseen predicted labels until reaching top_k.
+            else:
+                selected = set()
+                for y in ys:
+                    y = y.item()
+                    if y in pos:
+                        selected.add(y)
+                    if len(selected) == self.beam_width:
+                        break
+                if len(selected) < self.beam_width:
+                    selected = (list(selected) + list(pos - selected))[: self.beam_width]
+            train_clusters_selected[i] = np.asarray(list(selected))
+
+        train_labels_selected = PLTTrainer.cluster2label(clusters, train_clusters_selected)
+        val_labels_pred, val_scores_pred = PLTTrainer.cluster2label(clusters, val_clusters_pred, val_scores_pred)
+        num_labels_selected = self.beam_width * max(len(c) for c in clusters)
+
+        trainer = init_trainer(
+            self.checkpoint_dir,
+            epochs=self.epochs,
+            patience=self.patience,
+            early_stopping_metric=self.val_metric,
+            val_metric=self.val_metric,
+            silent=self.silent,
+            use_cpu=self.use_cpu,
+            limit_train_batches=self.limit_train_batches,
+            limit_val_batches=self.limit_val_batches,
+            limit_test_batches=self.limit_test_batches,
+            save_checkpoints=True,
+        )
+        trainer.checkpoint_callback.filename = f"{self.CHECKPOINT_NAME}1"
+
+        # train & val dataloaders for training
+        train_dataloader = self.dataloader(
+            PLTDataset(
+                train_x,
+                train_y,
+                num_classes=self.num_classes,
+                num_labels_selected=num_labels_selected,
+                labels_selected=train_labels_selected,
+            ),
+            shuffle=self.shuffle,
+        )
+        val_dataloader = self.dataloader(
+            PLTDataset(
+                val_x,
+                val_y,
+                num_classes=self.num_classes,
+                num_labels_selected=num_labels_selected,
+                labels_selected=val_labels_pred,
+                label_scores=val_scores_pred,
+            ),
+        )
+
+        try:
+            network = getattr(networks, "AttentionXML_1")(
+                embed_vecs=self.embed_vecs, num_classes=len(self.classes), **dict(self.network_config)
+            )
+        except Exception:
+            raise AttributeError("Failed to initialize AttentionXML")
+
+        model_1 = PLTModel(
+            classes=self.classes,
+            word_dict=self.word_dict,
+            network=network,
+            log_path=self.log_path,
+            learning_rate=self.learning_rate,
+            optimizer=self.optimizer,
+            momentum=self.momentum,
+            weight_decay=self.weight_decay,
+            lr_scheduler=self.lr_scheduler,
+            scheduler_config=self.scheduler_config,
+            val_metric=self.val_metric,
+            metric_threshold=self.metric_threshold,
+            monitor_metrics=self.monitor_metrics,
+            multiclass=self.multiclass,
+            loss_function=self.loss_function,
+            silent=self.silent,
+            save_k_predictions=self.save_k_predictions,
+        )
+        logger.info(f"Initialize model with weights from level 0")
+        # For weights not initialized by the level-0 model, use xavier uniform initialization
+        torch.nn.init.xavier_uniform_(model_1.network.attention.attention.weight)
+        # As the attention layer of model 1 is different from model 0, each layer needs to be initialized separately
+        model_1.network.embedding.load_state_dict(model_0.network.embedding.state_dict())
+        model_1.network.encoder.load_state_dict(model_0.network.encoder.state_dict())
+        model_1.network.output.load_state_dict(model_0.network.output.state_dict())
+
+        del model_0
+
+        logger.info(
+            f"Training level 1. Number of labels: {self.num_classes}."
+            f"Number of labels selected: {train_dataloader.dataset.num_labels_selected}"
+        )
+        trainer.fit(model_1, train_dataloader, val_dataloader)
+        logger.info(f"Best model loaded from {best_model_path}")
+        logger.info(f"Finish training level 1")
+
+    def test(self, dataset):
+        # retrieve word_dict from model_1
+        # prediction starts from level 0
+        model_0 = Model.load_from_checkpoint(
+            self.get_best_model_path(level=0),
+            save_k_predictions=self.beam_width,
+        )
+        model_1 = PLTModel.load_from_checkpoint(
+            self.get_best_model_path(level=1),
+            save_k_predictions=self.save_k_predictions,
+            metrics=self.metrics,
+        )
+        self.word_dict = model_1.word_dict
+        classes = model_1.classes
+
+        test_x = self.reformat_text(dataset)
+
+        if self.binarizer is None:
+            binarizer = MultiLabelBinarizer(classes=classes, sparse_output=True)
+            binarizer.fit(None)
+            test_y = binarizer.transform((i["label"] for i in dataset))
+        else:
+            test_y = self.binarizer.transform((i["label"] for i in dataset))
+        logger.info("Testing process started")
+        trainer = Trainer(
+            devices=1,
+            accelerator=self.accelerator,
+            logger=False,
+        )
+
+        test_dataloader = self.eval_dataloader(PlainDataset(test_x))
+
+        logger.info(f"Predicting level 0. Number of clusters: {self.beam_width}")
+        test_pred = trainer.predict(model_0, test_dataloader)
+        test_scores_pred = expit(np.vstack([i["top_k_pred_scores"] for i in test_pred]))
+        test_clusters_pred = np.vstack([i["top_k_pred"] for i in test_pred])
+
+        clusters = np.load(self.get_cluster_path(), allow_pickle=True)
+        test_labels_pred, test_scores_pred = PLTTrainer.cluster2label(clusters, test_clusters_pred, test_scores_pred)
+        num_labels_selected = self.beam_width * max(len(c) for c in clusters)
+
+        test_dataloader = self.eval_dataloader(
+            PLTDataset(
+                test_x,
+                test_y,
+                num_classes=self.num_classes,
+                num_labels_selected=num_labels_selected,
+                labels_selected=test_labels_pred,
+                label_scores=test_scores_pred,
+            ),
+        )
+
+        logger.info(f"Testing level 1")
+        trainer.test(model_1, test_dataloader)
+        logger.info("Testing process finished")
+
+        if self.save_k_predictions > 0:
+            batch_predictions = trainer.predict(model_1, test_dataloader)
+            pred_labels = np.vstack([batch["top_k_pred"] for batch in batch_predictions])
+            pred_scores = np.vstack([batch["top_k_pred_scores"] for batch in batch_predictions])
+            with open(self.predict_out_path, "w") as fp:
+                for pred_label, pred_score in zip(pred_labels, pred_scores):
+                    out_str = " ".join(
+                        [f"{model_1.classes[label]}:{score:.4}" for label, score in zip(pred_label, pred_score)]
+                    )
+                    fp.write(out_str + "\n")
+            logging.info(f"Saved predictions to: {self.predict_out_path}")
+
+        dump_log(self.log_path, config=self.config)
+
+    def reformat_text(self, dataset):
+        # Convert words to numbers according to their indices in word_dict. Then pad each instance to a certain length.
+        encoded_text = list(
+            map(
+                lambda text: torch.tensor([self.word_dict[word] for word in text], dtype=torch.int64)
+                if text
+                else torch.tensor([self.word_dict[UNK]], dtype=torch.int64),
+                [instance["text"][: self.max_seq_length] for instance in dataset],
+            )
+        )
+        # pad the first entry to be of length 500 if necessary
+        encoded_text[0] = torch.cat(
+            (
+                encoded_text[0],
+                torch.tensor(0, dtype=torch.int64).repeat(self.max_seq_length - encoded_text[0].shape[0]),
+            )
+        )
+        encoded_text = pad_sequence(encoded_text, batch_first=True)
+        return encoded_text
+
+    def get_best_model_path(self, level: int) -> Path:
+        return self.checkpoint_dir / f"{self.CHECKPOINT_NAME}{level}{ModelCheckpoint.FILE_EXTENSION}"
+
+    def get_cluster_path(self) -> Path:
+        return self.checkpoint_dir / f"{CLUSTER_NAME}{CLUSTER_FILE_EXTENSION}"
+
+
+###################################### Model ######################################
+
+
+class PLTModel(Model):
+    def __init__(
+        self,
+        classes,
+        word_dict,
+        network,
+        loss_function="binary_cross_entropy_with_logits",
+        log_path=None,
+        **kwargs,
+    ):
+        super().__init__(
+            classes=classes,
+            word_dict=word_dict,
+            network=network,
+            loss_function=loss_function,
+            log_path=log_path,
+            **kwargs,
+        )
+
+    def scatter_logits(
+        self,
+        logits: Tensor,
+        labels_selected: Tensor,
+        label_scores: Tensor,
+    ) -> Tensor:
+        """For each instance, we only have predictions on selected labels. This subroutine maps these predictions to
+        the whole label space. The scores of unsampled labels are set to 0."""
+        src = torch.sigmoid(logits.detach()) * label_scores
+        # During validation/testing, many fake labels might exist in a batch for the purpose of padding.
+        # A fake label has index len(classes) and does not belong to the real label space.
+        preds = torch.zeros(
+            labels_selected.size(0), len(self.classes) + 1, device=labels_selected.device, dtype=src.dtype
+        )
+        preds.scatter_(dim=1, index=labels_selected, src=src)
+        # slicing removes fake labels whose index is exactly len(self.classes)
+        # afterwards, preds is restored to the real label space
+        preds = preds[:, :-1]
+        return preds
+
+    def shared_step(self, batch):
+        """Return loss and predicted logits of the network.
+
+        Args:
+            batch (dict): A batch of text and label.
+
+        Returns:
+            loss (torch.Tensor): Loss between target and predict logits.
+            pred_logits (torch.Tensor): The predict logits (batch_size, num_classes).
+        """
+        y = torch.take_along_dim(batch["label"], batch["labels_selected"], dim=1)
+        pred_logits = self(batch)
+        loss = self.loss_function(pred_logits, y.float())
+        return loss, pred_logits
+
+    def _shared_eval_step(self, batch, batch_idx):
+        logits = self(batch)
+        logits = self.scatter_logits(logits, batch["labels_selected"], batch["label_scores"])
+        self.eval_metric.update(logits, batch["label"].long())
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        logits = self(batch)
+        scores, labels = torch.topk(torch.sigmoid(logits) * batch["label_scores"], self.save_k_predictions)
+        # This calculation is to align with LibMultiLabel class where logits rather than probabilities are returned
+        logits = torch.logit(scores)
+        return {
+            "top_k_pred": torch.take_along_dim(batch["labels_selected"], labels, dim=1).numpy(force=True),
+            "top_k_pred_scores": logits.numpy(force=True),
+        }
+
+
+###################################### Dataset ######################################
+
+
+class PlainDataset(Dataset):
+    """Plain (compared to nn.data_utils.TextDataset) dataset class for multi-label dataset.
+    WHY EXISTS: The reason why this class is necessary is that it can process labels in sparse format, while TextDataset
+    does not.
+    Moreover, TextDataset implements multilabel binarization in a mandatory way. Nevertheless, AttentionXML already does
+    this while generating clusters. There is no need to do multilabel binarization again.
+
+    Args:
+        x (list | ndarray | Tensor): texts.
+        y (Optional: csr_matrix | ndarray | Tensor): labels.
+    """
+
+    def __init__(self, x, y=None):
+        if y is not None:
+            assert len(x) == y.shape[0], "Sizes mismatch between texts and labels"
+        self.x = x
+        self.y = y
+
+    def __getitem__(self, idx: int) -> tuple[Sequence, ndarray] | tuple[Sequence]:
+        item = {"text": self.x[idx]}
+
+        # train/val/test
+        if self.y is not None:
+            if issparse(self.y):
+                y = self.y[idx].toarray().squeeze(0).astype(np.int32)
+            elif isinstance(self.y, ndarray):
+                y = self.y[idx].astype(np.int32)
+            elif is_tensor(self.y):
+                y = self.y[idx].int()
+            else:
+                raise TypeError(
+                    "The type of y should be one of scipy.csr_matrix, numpy.ndarry, and torch.Tensor."
+                    f"But got {type(self.y)} instead."
+                )
+            item["label"] = y
+        return item
+
+    def __len__(self):
+        return len(self.x)
+
+
+class PLTDataset(PlainDataset):
+    """Dataset for model_1 of AttentionXML.
+
+    Args:
+        x: texts.
+        y: labels.
+        num_classes: number of classes.
+        num_labels_selected: the number of selected labels.
+        labels_selected: sampled predicted labels from model_0. Shape: (len(x), predict_top_k).
+        label_scores: scores for each label. Shape: (len(x), predict_top_k).
+    """
+
+    def __init__(
+        self,
+        x,
+        y: Optional[csr_matrix | ndarray] = None,
+        *,
+        num_classes: int,
+        num_labels_selected: int,
+        labels_selected: ndarray | Tensor,
+        label_scores: Optional[ndarray | Tensor] = None,
+    ):
+        super().__init__(x, y)
+        self.num_classes = num_classes
+        self.num_labels_selected = num_labels_selected
+        self.labels_selected = labels_selected
+        self.label_scores = label_scores
+
+    def __getitem__(self, idx: int):
+        item = {"text": self.x[idx], "labels_selected": np.asarray(self.labels_selected[idx])}
+
+        if self.y is not None:
+            item["label"] = self.y[idx].toarray().squeeze(0).astype(np.int32)
+
+        # PyTorch requires inputs to be of the same shape. Pad any instance with length below num_labels_selected by
+        # randomly selecting labels.
+        # training
+        if self.label_scores is None:
+            # randomly add real labels when the number is below num_labels_selected
+            # some labels might be selected more than once
+            if len(item["labels_selected"]) < self.num_labels_selected:
+                samples = np.random.randint(
+                    self.num_classes,
+                    size=self.num_labels_selected - len(item["labels_selected"]),
+                )
+                item["labels_selected"] = np.concatenate([item["labels_selected"], samples])
+
+        # val/test/pred
+        else:
+            item["label_scores"] = self.label_scores[idx]
+            # add fake labels when the number of labels is below num_labels_selected
+            if len(item["labels_selected"]) < self.num_labels_selected:
+                item["label_scores"] = np.concatenate(
+                    [
+                        item["label_scores"],
+                        [-np.inf] * (self.num_labels_selected - len(item["labels_selected"])),
+                    ]
+                )
+                item["labels_selected"] = np.concatenate(
+                    [
+                        item["labels_selected"],
+                        [self.num_classes] * (self.num_labels_selected - len(item["labels_selected"])),
+                    ]
+                )
+        return item
+
+
+###################################### Cluster ######################################
+
+CLUSTER_NAME = "label_clusters"
+CLUSTER_FILE_EXTENSION = FILE_EXTENSION = ".npy"
+
+
+def build_label_tree(sparse_x: csr_matrix, sparse_y: csr_matrix, cluster_size: int, output_dir: str | Path):
+    """Build a binary tree to group labels into clusters, each of which contains up tp cluster_size labels. The tree has
+    several layers; nodes in the last layer correspond to the output clusters.
+    Given a set of labels (0, 1, 2, 3, 4, 5) and a cluster size of 2, the resulting clusters look something like:
+    ((0, 2), (1, 3), (4, 5)).
+
+    Args:
+        sparse_x: features extracted from texts in CSR sparse format
+        sparse_y: binarized labels in CSR sparse format
+        cluster_size: the maximum number of labels within each cluster
+        output_dir: directory to store the clustering file
+    """
+    # skip constructing label tree if the output file already exists
+    output_dir = output_dir if isinstance(output_dir, Path) else Path(output_dir)
+    cluster_path = output_dir / f"{CLUSTER_NAME}{FILE_EXTENSION}"
+    if cluster_path.exists():
+        logger.info("Clustering has finished in a previous run")
+        return
+
+    # meta info
+    logger.info("Label clustering started")
+    logger.info(f"Cluster size: {cluster_size}")
+    # The height of the tree satisfies the following inequality:
+    # 2**(tree_height - 1) * cluster_size < num_labels <= 2**tree_height * cluster_size
+    height = int(np.ceil(np.log2(sparse_y.shape[1] / cluster_size)))
+    logger.info(f"Labels will be grouped into {2 ** height} clusters")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # For each label, sum up normalized instances relevant to the label and normalize to get the label representation
+    label_repr = normalize(sparse_y.T @ csc_matrix(normalize(sparse_x)))
+
+    # clustering by a binary tree:
+    # at each layer split each cluster to two. Leave nodes correspond to the obtained clusters.
+    clusters = [np.arange(sparse_y.shape[1])]
+    for _ in range(height):
+        next_clusters = []
+        for cluster in clusters:
+            next_clusters += _split_cluster(cluster, label_repr[cluster])
+        clusters = next_clusters
+        logger.info(f"Having grouped {len(clusters)} clusters")
+
+    np.save(cluster_path, np.asarray(clusters, dtype=object))
+    logger.info(f"Label clustering finished. Saving results to {cluster_path}")
+
+
+def _split_cluster(cluster: ndarray, label_repr: csr_matrix) -> tuple[ndarray, ndarray]:
+    """A variant of KMeans implemented in AttentionXML. Here K = 2. The cluster is partitioned into two groups, each
+    with approximately equal size. Its main differences with the KMeans algorithm in scikit-learn are:
+    1. the distance metric is cosine similarity.
+    2. the end-of-loop criterion is the difference between the new and old average in-cluster distances to centroids.
+
+    Args:
+        cluster: a subset of labels
+        label_repr: the normalized representations of the relationship between labels and texts of the given cluster
+    """
+    # Randomly choose two points as initial centroids and obtain their label representations
+    centroids = label_repr[np.random.choice(len(cluster), size=2, replace=False)].toarray()
+
+    # Initialize distances (cosine similarity)
+    # Cosine similarity always falls to the interval [-1, 1]
+    old_dist = -2.0
+    new_dist = -1.0
+
+    # "c" denotes clusters
+    c0_idx = None
+    c1_idx = None
+
+    while new_dist - old_dist >= 1e-4:
+        # Notice that label_repr and centroids.T have been normalized
+        # Thus, dist indicates the cosine similarity between points and centroids.
+        dist = label_repr @ centroids.T  # shape: (n, 2)
+
+        # generate clusters
+        # let a = dist[:, 1] - dist[:, 0], the larger the element in a is, the closer the point is to c1
+        k = len(cluster) // 2
+        c_idx = np.argpartition(dist[:, 1] - dist[:, 0], kth=k)
+        c0_idx = c_idx[:k]
+        c1_idx = c_idx[k:]
+
+        # update distances
+        # the new distance is the average of in-cluster distances to the centroids
+        old_dist = new_dist
+        new_dist = (dist[c0_idx, 0].sum() + dist[c1_idx, 1].sum()) / len(cluster)
+
+        # update centroids
+        # the new centroid is the normalized average of the points in the cluster
+        centroids = normalize(
+            np.asarray(
+                [
+                    np.squeeze(np.asarray(label_repr[c0_idx].sum(axis=0))),
+                    np.squeeze(np.asarray(label_repr[c1_idx].sum(axis=0))),
+                ]
+            )
+        )
+    return cluster[c0_idx], cluster[c1_idx]

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/nn/model.py

@@ -155,7 +155,7 @@ class MultiLabelModel(L.LightningModule):
         Returns:
             dict: Top k label indexes and the prediction scores.
         """
-        _, pred_logits = self.shared_step(batch)
+        pred_logits = self(batch)
         pred_scores = pred_logits.detach().cpu().numpy()
         k = self.save_k_predictions
         top_k_idx = argsort_top_k(pred_scores, k, axis=1)
@@ -163,6 +163,10 @@ class MultiLabelModel(L.LightningModule):
 
         return {"top_k_pred": top_k_idx, "top_k_pred_scores": top_k_scores}
 
+    def forward(self, batch):
+        """compute predicted logits"""
+        return self.network(batch)["logits"]
+
     def print(self, *args, **kwargs):
         """Prints only from process 0 and not in silent mode. Use this in any
         distributed mode to log only once."""
@@ -178,7 +182,6 @@ class Model(MultiLabelModel):
     Args:
         classes (list): List of class names.
         word_dict (torchtext.vocab.Vocab): A vocab object which maps tokens to indices.
-        embed_vecs (torch.Tensor): The pre-trained word vectors of shape (vocab_size, embed_dim).
         network (nn.Module): Network (i.e., CAML, KimCNN, or XMLCNN).
         loss_function (str, optional): Loss function name (i.e., binary_cross_entropy_with_logits,
             cross_entropy). Defaults to 'binary_cross_entropy_with_logits'.
@@ -189,7 +192,6 @@ class Model(MultiLabelModel):
         self,
         classes,
         word_dict,
-        embed_vecs,
         network,
         loss_function="binary_cross_entropy_with_logits",
         log_path=None,
@@ -200,7 +202,6 @@ class Model(MultiLabelModel):
             ignore=["log_path"]
         )  # If log_path is saved, loading the checkpoint will cause an error since each experiment has unique log_path (result_dir).
         self.word_dict = word_dict
-        self.embed_vecs = embed_vecs
         self.classes = classes
         self.network = network
         self.configure_loss_function(loss_function)
@@ -224,8 +225,7 @@ class Model(MultiLabelModel):
             pred_logits (torch.Tensor): The predict logits (batch_size, num_classes).
         """
         target_labels = batch["label"]
-        outputs = self.network(batch)
-        pred_logits = outputs["logits"]
+        pred_logits = self(batch)
         loss = self.loss_function(pred_logits, target_labels.float())
 
         return loss, pred_logits

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/nn/networks/__init__.py

@@ -9,6 +9,7 @@ from .labelwise_attention_networks import BiGRULWAN
 from .labelwise_attention_networks import BiLSTMLWAN
 from .labelwise_attention_networks import BiLSTMLWMHAN
 from .labelwise_attention_networks import CNNLWAN
+from .labelwise_attention_networks import AttentionXML_0, AttentionXML_1
 
 
 def get_init_weight_func(init_weight):

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/nn/networks/labelwise_attention_networks.py

@@ -10,6 +10,8 @@ from .modules import (
     LabelwiseAttention,
     LabelwiseMultiHeadAttention,
     LabelwiseLinearOutput,
+    PartialLabelwiseAttention,
+    MultilayerLinearOutput,
 )
 
 
@@ -266,3 +268,70 @@ class CNNLWAN(LabelwiseAttentionNetwork):
         x, _ = self.attention(x)  # (batch_size, num_classes, hidden_dim)
         x = self.output(x)  # (batch_size, num_classes)
         return {"logits": x}
+
+
+class AttentionXML_0(nn.Module):
+    def __init__(
+        self,
+        embed_vecs,
+        num_classes: int,
+        rnn_dim: int,
+        linear_size: list,
+        freeze_embed_training: bool = False,
+        rnn_layers: int = 1,
+        embed_dropout: float = 0.2,
+        encoder_dropout: float = 0,
+        post_encoder_dropout: float = 0.5,
+    ):
+        super().__init__()
+        self.embedding = Embedding(embed_vecs, freeze=freeze_embed_training, dropout=embed_dropout)
+        self.encoder = LSTMEncoder(embed_vecs.shape[1], rnn_dim // 2, rnn_layers, encoder_dropout, post_encoder_dropout)
+        self.attention = LabelwiseAttention(rnn_dim, num_classes)
+        self.output = MultilayerLinearOutput([rnn_dim] + linear_size, 1)
+
+    def forward(self, inputs):
+        x = inputs["text"]
+        # the index of padding is 0
+        masks = x != 0
+        lengths = masks.sum(dim=1)
+        masks = masks[:, : lengths.max()]
+
+        x = self.embedding(x)[:, : lengths.max()]  # batch_size, length, embedding_size
+        x = self.encoder(x, lengths)  # batch_size, length, hidden_size
+        x, _ = self.attention(x)  # batch_size, num_classes, hidden_size
+        x = self.output(x)  # batch_size, num_classes
+        return {"logits": x}
+
+
+class AttentionXML_1(nn.Module):
+    def __init__(
+        self,
+        embed_vecs,
+        num_classes: int,
+        rnn_dim: int,
+        linear_size: list,
+        freeze_embed_training: bool = False,
+        rnn_layers: int = 1,
+        embed_dropout: float = 0.2,
+        encoder_dropout: float = 0,
+        post_encoder_dropout: float = 0.5,
+    ):
+        super().__init__()
+        self.embedding = Embedding(embed_vecs, freeze=freeze_embed_training, dropout=embed_dropout)
+        self.encoder = LSTMEncoder(embed_vecs.shape[1], rnn_dim // 2, rnn_layers, encoder_dropout, post_encoder_dropout)
+        self.attention = PartialLabelwiseAttention(rnn_dim, num_classes)
+        self.output = MultilayerLinearOutput([rnn_dim] + linear_size, 1)
+
+    def forward(self, inputs):
+        x = inputs["text"]
+        labels_selected = inputs["labels_selected"]
+        # the index of padding is 0
+        masks = x != 0
+        lengths = masks.sum(dim=1)
+        masks = masks[:, : lengths.max()]
+
+        x = self.embedding(x)[:, : lengths.max()]  # batch_size, length, embedding_size
+        x = self.encoder(x, lengths)  # batch_size, length, hidden_size
+        x, _ = self.attention(x, labels_selected)  # batch_size, sample_size, hidden_size
+        x = self.output(x)  # batch_size, sample_size
+        return {"logits": x}

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/nn/networks/modules.py

@@ -210,3 +210,35 @@ class LabelwiseLinearOutput(nn.Module):
 
     def forward(self, input):
         return (self.output.weight * input).sum(dim=-1) + self.output.bias
+
+
+class PartialLabelwiseAttention(nn.Module):
+    """Similar to LabelwiseAttention.
+    What makes the class different from LabelwiseAttention is that only the weights of selected labels will be
+    updated in a single iteration.
+    """
+
+    def __init__(self, hidden_size, num_classes):
+        super().__init__()
+        self.attention = nn.Embedding(num_classes + 1, hidden_size)
+
+    def forward(self, inputs, labels_selected):
+        attn_inputs = inputs.transpose(1, 2)  # batch_size, hidden_dim, length
+        attn_weights = self.attention(labels_selected)  # batch_size, sample_size, hidden_dim
+        attention = attn_weights @ attn_inputs  # batch_size, sample_size, length
+        attention = F.softmax(attention, -1)  # batch_size, sample_size, length
+        logits = attention @ inputs  # batch_size, sample_size, hidden_dim
+        return logits, attention
+
+
+class MultilayerLinearOutput(nn.Module):
+    def __init__(self, linear_size, output_size):
+        super().__init__()
+        self.linears = nn.ModuleList(nn.Linear(in_s, out_s) for in_s, out_s in zip(linear_size[:-1], linear_size[1:]))
+        self.output = nn.Linear(linear_size[-1], output_size)
+
+    def forward(self, inputs):
+        linear_out = inputs
+        for linear in self.linears:
+            linear_out = F.relu(linear(linear_out))
+        return torch.squeeze(self.output(linear_out), -1)

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel/nn/nn_utils.py

@@ -100,7 +100,6 @@ def init_model(
     model = Model(
         classes=classes,
         word_dict=word_dict,
-        embed_vecs=embed_vecs,
         network=network,
         log_path=log_path,
         learning_rate=learning_rate,

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: libmultilabel
-Version: 0.6.2
-Summary: A library for multi-label text classification
+Version: 0.7.0
+Summary: A library for multi-class and multi-label classification
 Home-page: https://github.com/ASUS-AICS/LibMultiLabel
 Author: LibMultiLabel Team
 License: MIT License

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/libmultilabel.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ libmultilabel/linear/preprocessor.py
 libmultilabel/linear/tree.py
 libmultilabel/linear/utils.py
 libmultilabel/nn/__init__.py
+libmultilabel/nn/attentionxml.py
 libmultilabel/nn/data_utils.py
 libmultilabel/nn/metrics.py
 libmultilabel/nn/model.py

{libmultilabel-0.6.2 → libmultilabel-0.7.0}/setup.cfg

@@ -1,10 +1,10 @@
 [metadata]
 name = libmultilabel
-version = 0.6.2
+version = 0.7.0
 author = LibMultiLabel Team
 license = MIT License
 license_file = LICENSE
-description = A library for multi-label text classification
+description = A library for multi-class and multi-label classification
 long_description = See documentation here: https://www.csie.ntu.edu.tw/~cjlin/libmultilabel
 url = https://github.com/ASUS-AICS/LibMultiLabel
 project_urls =