Add argument to tempo_aug (#105)

EleutherAI · Mar 5, 2024 · 7c2c47e · 7c2c47e
1 parent 56b66ea
commit 7c2c47e
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/aria/tokenizer/tokenizer.py b/aria/tokenizer/tokenizer.py
@@ -819,16 +819,18 @@ def tempo_aug(
             instruments_wd: list,
             _tempo_aug_range: float,
             _mixup: bool,
+            tempo_aug: float | None = None,
         ):
             """This must be used with export_aug_fn_concat in order to work
             properly for concatenated sequences."""
 
             def _quantize_time(_n: int):
                 return round(_n / time_step) * time_step
 
-            tempo_aug = random.uniform(
-                1 - _tempo_aug_range, 1 + _tempo_aug_range
-            )
+            if not tempo_aug:
+                tempo_aug = random.uniform(
+                    1 - _tempo_aug_range, 1 + _tempo_aug_range
+                )
 
             src_time_tok_cnt = 0
             dim_tok_seen = None