Update train.py

train.py

@@ -31,7 +31,6 @@ class ASR(sb.Brain):
         # Add waveform augmentation if specified.
         if stage == sb.Stage.TRAIN:
             sig, self.sig_lens = self.hparams.wav_augment(sig, self.sig_lens)
-            # tokens_bos = self.hparams.wav_augment.replicate_labels(tokens_bos)
 
         # Forward pass
         encoded_outputs = self.modules.encoder_w2v2(sig.detach())
@@ -75,9 +74,6 @@ class ASR(sb.Brain):
         tokens_eos, tokens_eos_lens = batch.tokens_eos
         tokens, tokens_lens = batch.tokens
 
-        # if stage == sb.Stage.TRAIN:
-            # (tokens, tokens_lens, tokens_eos, tokens_eos_lens) = self.hparams.wav_augment.replicate_multiple_labels(tokens, tokens_lens, tokens_eos, tokens_eos_lens)
-
         loss = self.hparams.nll_cost(log_probabilities=predictions["seq_logprobs"], targets=tokens_eos, length=tokens_eos_lens)
 
         if self.is_ctc_active(stage):
@@ -87,8 +83,6 @@ class ASR(sb.Brain):
             loss += self.hparams.ctc_weight * loss_ctc
 
         if stage != sb.Stage.TRAIN:
-            # for prediction in predictions["tokens"]:
-                # print(self.hparams.tokenizer.decode_ids(prediction))
             predicted_words = [self.hparams.tokenizer.decode_ids(prediction).split(" ") for prediction in predictions["tokens"]]
             target_words = [words.split(" ") for words in batch.transcript]
             self.wer_metric.append(ids, predicted_words, target_words)
@@ -118,7 +112,6 @@ class ASR(sb.Brain):
             sb.nnet.schedulers.update_learning_rate(self.optimizer, new_lr)
             self.hparams.train_logger.log_stats(
                 stats_meta={"epoch": epoch, "lr": old_lr},
-                # stats_meta={"epoch": epoch},
                 train_stats=self.train_stats,
                 valid_stats=stage_stats,
             )
@@ -155,7 +148,6 @@ class ASR(sb.Brain):
         with torch.no_grad():
             true_labels = []
             pred_labels = []
-            #for batch in tqdm(dataset, dynamic_ncols=True):
             for batch in dataset:
                 # Make sure that your compute_forward returns the predictions !!!
                 # In the case of the template, when stage = TEST, a beam search is applied 
@@ -167,39 +159,84 @@ class ASR(sb.Brain):
 
                 predicted_words = [self.hparams.tokenizer.decode_ids(prediction).split(" ") for prediction in predictions["tokens"]]
                 for sent in predicted_words:
+                    # sent = " ".join(sent)
+                    sent = filter_repetitions(sent, 3)
                     sent = " ".join(sent)
                     pred_batch.append(sent)
 
                 pred_labels.append(pred_batch[0])
                 true_labels.append(batch.transcript[0])
 
-                # print("True: ", batch.transcript[0])
-                # print("Pred: ", pred_batch[0])
-                # with open("predictions/predictions_arhiv.txt", "a") as f:
-                    # f.write("True: " + batch.transcript[0] + "\n")
-                    # f.write("Pred: " + pred_batch[0] + "\n\n")
-                print("True: ", batch.transcript[0])
-                print("Pred: ", pred_batch[0])
-            
         print('WER: ', wer(true_labels, pred_labels) * 100)
         print('CER: ', cer(true_labels, pred_labels) * 100)
 
 
+def filter_repetitions(seq, max_repetition_length):
+        seq = list(seq)
+        output = []
+        max_n = len(seq) // 2
+        for n in range(max_n, 0, -1):
+            max_repetitions = max(max_repetition_length // n, 1)
+            # Don't need to iterate over impossible n values:
+            # len(seq) can change a lot during iteration
+            if (len(seq) <= n*2) or (len(seq) <= max_repetition_length):
+                continue
+            iterator = enumerate(seq)
+            # Fill first buffers:
+            buffers = [[next(iterator)[1]] for _ in range(n)]
+            for seq_index, token in iterator:
+                current_buffer = seq_index % n
+                if token != buffers[current_buffer][-1]:
+                    # No repeat, we can flush some tokens
+                    buf_len = sum(map(len, buffers))
+                    flush_start = (current_buffer-buf_len) % n
+                    # Keep n-1 tokens, but possibly mark some for removal
+                    for flush_index in range(buf_len - buf_len%n):
+                        if (buf_len - flush_index) > n-1:
+                            to_flush = buffers[(flush_index + flush_start) % n].pop(0)
+                        else:
+                            to_flush = None
+                        # Here, repetitions get removed:
+                        if (flush_index // n < max_repetitions) and to_flush is not None:
+                            output.append(to_flush)
+                        elif (flush_index // n >= max_repetitions) and to_flush is None:
+                            output.append(to_flush)
+                buffers[current_buffer].append(token)
+            # At the end, final flush
+            current_buffer += 1
+            buf_len = sum(map(len, buffers))
+            flush_start = (current_buffer-buf_len) % n
+            for flush_index in range(buf_len):
+                to_flush = buffers[(flush_index + flush_start) % n].pop(0)
+                # Here, repetitions just get removed:
+                if flush_index // n < max_repetitions:
+                    output.append(to_flush)
+            seq = []
+            to_delete = 0
+            for token in output:
+                if token is None:
+                    to_delete += 1
+                elif to_delete > 0:
+                    to_delete -= 1
+                else:
+                    seq.append(token)
+            output = []
+        return seq
+
 def dataio_prepare(hparams):
     """This function prepares the datasets to be used in the brain class.
     It also defines the data processing pipeline through user-defined functions.
     """
     data_folder = hparams["data_folder"]
 
-    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(json_path=os.path.join(hparams["data_folder"], "train_corrected.json"), replacements={"data_root": data_folder})
-
+    train_data = sb.dataio.dataset.DynamicItemDataset.from_json(json_path=os.path.join(hparams["data_folder"], "train.json"), replacements={"data_root": data_folder})
     train_data = train_data.filtered_sorted(sort_key="duration")
     hparams["train_dataloader_opts"]["shuffle"] = False
 
-    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(json_path=os.path.join(hparams["data_folder"], "dev_corrected.json"), replacements={"data_root": data_folder})
+    valid_data = sb.dataio.dataset.DynamicItemDataset.from_json(json_path=os.path.join(hparams["data_folder"], "dev.json"), replacements={"data_root": data_folder})
     valid_data = valid_data.filtered_sorted(sort_key="duration")
 
-    test_data = sb.dataio.dataset.DynamicItemDataset.from_json(json_path=os.path.join(hparams["data_folder"], "test_arhiv.json"), replacements={"data_root": data_folder})
+    test_data = sb.dataio.dataset.DynamicItemDataset.from_json(json_path=os.path.join(hparams["data_folder"], "test.json"), replacements={"data_root": data_folder})
 
 
     datasets = [train_data, valid_data, test_data]
@@ -212,18 +249,8 @@ def dataio_prepare(hparams):
     @sb.utils.data_pipeline.takes("data_path")
     @sb.utils.data_pipeline.provides("sig")
     def audio_pipeline(data_path):
-        if "cv-mk" in data_path:
-            filename = data_path.split("clips")[1]
-            data_path = "/m/triton/scratch/elec/t405-puhe/p/porjazd1/macedonian_asr/data/CV-18_MK/cv-mk/mk/clips" + filename
-        elif "podcast" in data_path:
-            filename = data_path.split("segmented_audio")[1]
-            data_path = "/m/triton/scratch/elec/t405-puhe/p/porjazd1/macedonian_asr/data/podcast/audio/segmented_audio" + filename
-        elif "arhiv" in data_path:
-            filename = data_path.split("segmented_audio")[1]
-            data_path = "/m/triton/scratch/elec/t405-puhe/p/porjazd1/macedonian_asr/data/arhiv/audio/segmented_audio" + filename
-
         sig, sr = librosa.load(data_path, sr=16000)
-        # sig = sb.dataio.dataio.read_audio(wav)
+        # sig = sb.dataio.dataio.read_audio(wav) # alternatively use the SpeechBrain data loading function
         return sig
 
     sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline)
@@ -270,8 +297,6 @@ if __name__ == "__main__":
     # here we create the datasets objects as well as tokenization and encoding
     (train_data, valid_data, test_data) = dataio_prepare(hparams)
 
-    #run_on_main(hparams["pretrainer"].collect_files)
-    #hparams["pretrainer"].load_collected()
 
     # Trainer initialization
     asr_brain = ASR(