Version 0.1.3

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.1.3] 2023-11-20
+### Added
+- Forbid repetition mode argument to LightningModule and HuggingFace wrapper.
+
 ## [0.1.2] 2023-11-17
 ### Fixed
 - Task embeddings inputs `wavcaps_audioset_sl` and `wavcaps_bbc_sound_effects`.

src/conette/__init__.py

@@ -10,7 +10,7 @@
 __license__ = "MIT"
 __maintainer__ = "Etienne Labbé (Labbeti)"
 __status__ = "Development"
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 
 from pathlib import Path

src/conette/huggingface/model.py

@@ -167,6 +167,7 @@ def forward(
         beam_size: Optional[int] = None,
         min_pred_size: Optional[int] = None,
         max_pred_size: Optional[int] = None,
+        forbid_rep_mode: Optional[str] = None,
     ) -> dict[str, Any]:
         # Preprocessing (load data + encode features)
         if preprocess:
@@ -215,6 +216,7 @@ def forward(
             beam_size=beam_size,
             min_pred_size=min_pred_size,
             max_pred_size=max_pred_size,
+            forbid_rep_mode=forbid_rep_mode,
         )
         kwds = {k: v for k, v in kwds.items() if v is not None}
         outs = self.model(batch, **kwds)
@@ -233,6 +235,7 @@ def __call__(
         beam_size: Optional[int] = None,
         min_pred_size: Optional[int] = None,
         max_pred_size: Optional[int] = None,
+        forbid_rep_mode: Optional[str] = None,
     ) -> dict[str, Any]:
         return super().__call__(
             x=x,
@@ -242,4 +245,5 @@ def __call__(
             beam_size=beam_size,
             min_pred_size=min_pred_size,
             max_pred_size=max_pred_size,
+            forbid_rep_mode=forbid_rep_mode,
         )

src/conette/pl_modules/common.py

@@ -259,9 +259,8 @@ def get_forbid_rep_mask(
     else:
         FORBID_MODES = (
             "none",
+            "all",
             "content_words",
-            "content_words_not_bpe_part",
-            "content_words_stem_lemm",
         )
         raise ValueError(
             f"Invalid argument {forbid_rep_mode=}. (expected one of {FORBID_MODES})"

src/conette/pl_modules/conette.py

@@ -411,14 +411,22 @@ def decode_audio(
             if "bos_id" not in kwargs:
                 raise ValueError(f"Missing argument 'bos_id' for {decode_method=}.")
 
+            forbid_rep_mode = kwargs.get("forbid_rep_mode", None)
+            if forbid_rep_mode is None:
+                forbid_rep_mask = self.forbid_rep_mask
+            else:
+                forbid_rep_mask = get_forbid_rep_mask(
+                    forbid_rep_mode, self.tokenizer, self.device, self.hp.verbose
+                )
+
             generate_expt_hp = {
                 "pad_id": self.pad_id,
                 "bos_id": self.bos_id,
                 "eos_id": self.eos_id,
                 "vocab_size": self.tokenizer.get_vocab_size(),
                 "min_pred_size": self.hp.min_pred_size,
                 "max_pred_size": self.hp.max_pred_size,
-                "forbid_rep_mask": self.forbid_rep_mask,
+                "forbid_rep_mask": forbid_rep_mask,
                 "beam_size": self.hp.beam_size,
             }
             kwargs = generate_expt_hp | kwargs

src/conette/transforms/mixup.py

@@ -10,8 +10,7 @@
 
 
 def pann_mixup(x: Tensor, mixup_lambda: Tensor) -> Tensor:
-    """Mixup x of even indexes (0, 2, 4, ...) with x of odd indexes
-    (1, 3, 5, ...).
+    """Mixup x of even indexes (0, 2, 4, ...) with x of odd indexes (1, 3, 5, ...).
 
     Args:
       x: (batch_size * 2, ...)
@@ -28,7 +27,9 @@ def pann_mixup(x: Tensor, mixup_lambda: Tensor) -> Tensor:
 
 
 def sample_lambda(
-    alpha: Union[float, Tensor], asymmetric: bool, size: Iterable[int] = ()
+    alpha: Union[float, Tensor],
+    asymmetric: bool,
+    size: Iterable[int] = (),
 ) -> Tensor:
     """
     :param alpha: alpha hp to control the Beta distribution.