genhpf 1.0.11__py3-none-any.whl

genhpf/criterions/wav2vec2_criterion.py

@@ -0,0 +1,130 @@
+import math
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+import genhpf.utils.utils as utils
+from genhpf.criterions import BaseCriterion, register_criterion
+from genhpf.criterions.criterion import CriterionConfig
+from genhpf.loggings import metrics
+from genhpf.loggings.meters import safe_round
+
+
+@dataclass
+class Wav2Vec2CriterionConfig(CriterionConfig):
+    loss_weights: Optional[List[float]] = field(
+        default=None, metadata={"help": "weights for additional loss terms (not first one)"}
+    )
+
+
+@register_criterion("wav2vec2_criterion", dataclass=Wav2Vec2CriterionConfig)
+class Wav2Vec2Criterion(BaseCriterion):
+    def __init__(self, cfg: Wav2Vec2CriterionConfig):
+        super().__init__(cfg)
+
+        self.loss_weights = cfg.loss_weights
+
+    def compute_loss(
+        self, logits: torch.Tensor, targets: torch.Tensor, sample=None, net_output=None, model=None
+    ) -> Tuple[torch.Tensor, List[float]]:
+        """Compute the loss given the logits and targets from the model."""
+
+        losses = []
+
+        loss = F.cross_entropy(logits, targets, reduction="sum")
+        losses.append(loss.detach().item())
+
+        sample_size = self.get_sample_size(sample, targets)
+
+        if self.loss_weights is not None:
+            assert hasattr(model, "get_extra_losses")
+            extra_losses = model.get_extra_losses(net_output)
+            if torch.is_tensor(extra_losses):
+                extra_losses = [extra_losses]
+            if len(self.loss_weights) == 1 and len(extra_losses) != 1:
+                self.loss_weights = [self.weights[0]] * len(extra_losses)
+            assert len(extra_losses) == len(
+                self.loss_weights
+            ), f"{len(extra_losses)}, {len(self.loss_weights)}"
+            for p, coef in zip(extra_losses, self.loss_weights):
+                if coef != 0 and p is not None:
+                    p = coef * p.float() * sample_size
+                    loss += p
+                    losses.append(p.detach().item())
+
+        return loss, losses
+
+    def get_sample_size(self, sample, targets):
+        """
+        Get the sample size, which is used as the denominator for the gradient
+        """
+        if "sample_size" in sample:
+            sample_size = sample["sample_size"]
+        else:
+            sample_size = targets.numel()
+        return sample_size
+
+    def get_logging_outputs(
+        self, logging_output, logits: torch.Tensor, targets: torch.Tensor, sample=None
+    ) -> List[Dict[str, Any]]:
+        """
+        Get the logging output to display while training
+        """
+        with torch.no_grad():
+            if logits.numel() == 0:
+                corr = 0
+                count = 0
+            else:
+                assert logits.dim() > 1, logits.shape
+                max = logits.argmax(-1) == 0
+                min = logits.argmin(-1) == 0
+
+                both = max & min
+                corr = max.long().sum().item() - both.long().sum().item()
+                count = float(max.numel())
+
+            logging_output["correct"] = corr
+            logging_output["count"] = count
+        return logging_output
+
+    @staticmethod
+    def reduce_metrics(logging_outputs: List[Dict[str, Any]], prefix: str = None) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        if prefix is None:
+            prefix = ""
+        elif prefix is not None and not prefix.endswith("_"):
+            prefix = prefix + "_"
+
+        loss_sum = utils.item(sum(log.get("loss", 0) for log in logging_outputs))
+
+        sample_size = utils.item(sum(log.get("sample_size", 0) for log in logging_outputs))
+
+        metrics.log_scalar(f"{prefix}loss", loss_sum / (sample_size or 1) / math.log(2), sample_size, round=3)
+
+        correct = sum(log.get("correct", 0) for log in logging_outputs)
+        metrics.log_scalar(f"_{prefix}correct", correct)
+
+        total = sum(log.get("count", 0) for log in logging_outputs)
+        metrics.log_scalar(f"_{prefix}total", total)
+
+        if total > 0:
+            metrics.log_derived(
+                f"{prefix}accuracy",
+                lambda meters: safe_round(meters[f"_{prefix}correct"].sum / meters[f"_{prefix}total"].sum, 5)
+                if meters[f"_{prefix}total"].sum > 0
+                else float("nan"),
+            )
+
+        builtin_keys = {"loss", "sample_size", "correct", "count"}
+
+        for k in logging_outputs[0]:
+            if k not in builtin_keys:
+                val = sum(log.get(k, 0) for log in logging_outputs)
+                if k.startswith("loss"):
+                    metrics.log_scalar(
+                        prefix + k, val / (sample_size or 1) / math.log(2), sample_size, round=3
+                    )
+                else:
+                    metrics.log_scalar(prefix + k, val / len(logging_outputs), round=3)

genhpf/datasets/__init__.py

@@ -0,0 +1,84 @@
+from typing import List
+
+from genhpf.configs import Config
+
+from .dataset import BaseDataset
+from .genhpf_dataset import FlattenedGenHPFDataset, HierarchicalGenHPFDataset
+from .meds_dataset import HierarchicalMEDSDataset
+
+__all__ = ["BaseDataset", "HierarchicalGenHPFDataset", "FlattenedGenHPFDataset", "HierarchicalMEDSDataset"]
+
+
+def load_dataset(
+    data_path: str,
+    subsets: List[str],
+    cfg: Config,
+):
+    dataset_cfg = cfg.dataset
+    model_cfg = cfg.model
+    criterion_cfg = cfg.criterion
+
+    manifest_paths = [f"{data_path}/{subset.strip()}.tsv" for subset in subsets]
+
+    if dataset_cfg.data_format == "genhpf":
+        if model_cfg.structure == "hierarchical":
+            dataset = HierarchicalGenHPFDataset(
+                manifest_paths=manifest_paths,
+                label=dataset_cfg.label,
+                tasks=getattr(criterion_cfg, "task_names", None),
+                num_labels=getattr(criterion_cfg, "num_labels", None),
+                vocab_size=dataset_cfg.vocab_size,
+                pad_token_id=dataset_cfg.pad_token_id,
+                sep_token_id=dataset_cfg.sep_token_id,
+                ignore_index=dataset_cfg.ignore_index,
+                apply_mask=dataset_cfg.apply_mask or "mlm" in model_cfg._name,
+                mask_token_id=dataset_cfg.mask_token_id,
+                mask_prob=dataset_cfg.mask_prob,
+                mask_unit=dataset_cfg.mask_unit,
+                simclr="simclr" in model_cfg._name,
+                dummy_token_id=dataset_cfg.dummy_token_id,
+            )
+        else:
+            dataset = FlattenedGenHPFDataset(
+                manifest_paths=manifest_paths,
+                label=dataset_cfg.label,
+                tasks=getattr(criterion_cfg, "task_names", None),
+                num_labels=getattr(criterion_cfg, "num_labels", None),
+                vocab_size=dataset_cfg.vocab_size,
+                pad_token_id=dataset_cfg.pad_token_id,
+                sep_token_id=dataset_cfg.sep_token_id,
+                ignore_index=dataset_cfg.ignore_index,
+                apply_mask=dataset_cfg.apply_mask or "mlm" in model_cfg._name,
+                mask_token_id=dataset_cfg.mask_token_id,
+                mask_prob=dataset_cfg.mask_prob,
+                mask_unit=dataset_cfg.mask_unit,
+                simclr="simclr" in model_cfg._name,
+            )
+    elif dataset_cfg.data_format == "meds":
+        assert model_cfg.structure == "hierarchical", (
+            "we currently only support hierarchical structure for MEDS dataset."
+            " please set model.structure to 'hierarchical'"
+        )
+        dataset = HierarchicalMEDSDataset(
+            manifest_paths=manifest_paths,
+            max_events=model_cfg.agg_max_seq_len,
+            label=dataset_cfg.label,
+            tasks=getattr(criterion_cfg, "task_names", None),
+            num_labels=getattr(criterion_cfg, "num_labels", None),
+            structure=model_cfg.structure,
+            vocab_size=dataset_cfg.vocab_size,
+            pad_token_id=dataset_cfg.pad_token_id,
+            sep_token_id=dataset_cfg.sep_token_id,
+            ignore_index=dataset_cfg.ignore_index,
+            apply_mask=dataset_cfg.apply_mask or "mlm" in model_cfg._name,
+            mask_token_id=dataset_cfg.mask_token_id,
+            mask_prob=dataset_cfg.mask_prob,
+            mask_unit=dataset_cfg.mask_unit,
+            simclr="simclr" in model_cfg._name,
+            dummy_token_id=dataset_cfg.dummy_token_id,
+            debug=cfg.common.debug,
+        )
+    else:
+        raise NotImplementedError(f"unsupported data format: {dataset_cfg.data_format}")
+
+    return dataset

genhpf/datasets/dataset.py

@@ -0,0 +1,109 @@
+import logging
+from typing import Union
+
+import numpy as np
+import torch.utils.data
+
+from genhpf.configs import ChoiceEnum
+
+MASK_UNIT_CHOICES = ChoiceEnum(["token", "event"])
+
+logger = logging.getLogger(__name__)
+
+
+class BaseDataset(torch.utils.data.Dataset):
+    def __len__(self):
+        raise NotImplementedError
+
+    def __getitem__(self, index):
+        raise NotImplementedError
+
+    def collator(self, samples):
+        raise NotImplementedError
+
+    def mask(
+        self,
+        tokens: Union[np.ndarray, torch.Tensor],
+        mask_prob: float,
+        vocab_size: int,
+        mask_token_id: int,
+        mask_unit: MASK_UNIT_CHOICES = "token",
+        sep_token_id: int = None,
+    ):
+        """
+        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
+        """
+        assert 0.0 < mask_prob < 1.0, "mask_prob must be in the range (0.0, 1.0)"
+
+        if isinstance(tokens, np.ndarray):
+            tokens = torch.LongTensor(tokens)
+        tokens = tokens.long()
+        labels = tokens.clone()
+
+        assert tokens.dim() == 2, (
+            "input tokens must be 2D tensor, where the first dimension is the number of embeddings "
+            "(i.e., input_ids, type_ids, dpe_ids) and the second dimension is the length of token "
+            "sequence."
+        )
+
+        if mask_unit == "token":
+            probability_matrix = torch.full((labels.size(-1),), mask_prob)
+            # do not mask special tokens
+            if not hasattr(self, "tokenizer"):
+                from transformers import AutoTokenizer
+
+                self.tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+            special_tokens_mask = torch.tensor(
+                self.tokenizer.get_special_tokens_mask(labels[0], already_has_special_tokens=True),
+                dtype=torch.bool,
+            )
+            probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
+            probability_matrix[torch.where(labels[0] == self.pad_token_id)] = 0.0
+
+            mask_indices = torch.bernoulli(probability_matrix).bool()
+            while mask_indices.sum() == 0:
+                mask_indices = torch.bernoulli(probability_matrix).bool()
+        elif mask_unit == "event":
+            if sep_token_id is None:
+                logger.warning(
+                    "sep_token_id is not provided. Using the default [SEP] token id (102) as "
+                    "the event delimiter."
+                )
+                sep_token_id = 102  # token id for [SEP]
+            event_indices = torch.where(tokens[0] == sep_token_id)[0]
+            assert len(event_indices) > 1, (
+                "there must be at least two events in the input sequence to apply span masking. "
+                "check if you are using the hierarchical structure which is not supporting the "
+                "span masking method."
+            )
+            mask_indices = torch.zeros_like(tokens[0]).bool()
+            masked_event_indices = torch.randperm(len(event_indices) - 1)[
+                : round((len(event_indices) - 1) * mask_prob)
+            ]
+            for i in masked_event_indices:
+                mask_indices[event_indices[i] : event_indices[i + 1]] = True
+        else:
+            raise ValueError(f"mask_unit must be one of {MASK_UNIT_CHOICES}")
+
+        labels[:, ~mask_indices] = self.ignore_index  # we only compute loss on masked tokens
+
+        # 80% of the time, we replace the masked input tokens with tokenizer.mask_token ([MASK])
+        mask_positions = torch.bernoulli(torch.full(mask_indices.shape, 0.8)).bool() & mask_indices
+        tokens[0, mask_positions] = mask_token_id  # for input_ids
+        tokens[1, mask_positions] = 4  # for type_ids
+        tokens[2, mask_positions] = 15  # for dpe_ids
+
+        # 10% of the time, we replace the masked input tokens with random word
+        random_positions = (
+            torch.bernoulli(torch.full(mask_indices.shape, 0.5)).bool() & mask_indices & ~mask_positions
+        )
+        random_words = torch.randint(vocab_size, mask_indices.shape, dtype=torch.long)
+        tokens[0, random_positions] = random_words[random_positions]
+        random_types = torch.randint(7, mask_indices.shape, dtype=torch.long)
+        tokens[1, random_positions] = random_types[random_positions]
+        random_dpes = torch.randint(16, mask_indices.shape, dtype=torch.long)
+        tokens[2, random_positions] = random_dpes[random_positions]
+
+        # the rest of the time, we keep the masked input tokens unchanged
+        return tokens, labels