first upload

README.md

@@ -1,5 +1,5 @@
 ---
-title: Artst
+title: ArtstTTS
 emoji: 🔥
 colorFrom: yellow
 colorTo: gray
@@ -7,6 +7,7 @@ sdk: gradio
 sdk_version: 4.7.1
 app_file: app.py
 pinned: false
+python_version: 3.8.2
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,59 @@
+import os
+import torch
+import gradio as gr
+import os.path as op
+import pyarabic.araby as araby
+
+from artst.tasks.artst import ArTSTTask
+from transformers import SpeechT5HifiGan
+from artst.models.artst import ArTSTTransformerModel
+from fairseq.tasks.hubert_pretraining import LabelEncoder
+from fairseq.data.audio.speech_to_text_dataset import get_features_or_waveform 
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+WORK_DIR = os.getcwd()
+checkpoint = torch.load('ckpts/clartts_tts.pt')
+checkpoint['cfg']['task'].t5_task = 't2s'
+task = ArTSTTask.setup_task(checkpoint['cfg']['task'])
+
+emb_path='embs/clartts.npy'
+model = ArTSTTransformerModel.build_model(checkpoint['cfg']['model'], task)
+model.load_state_dict(checkpoint['model'])
+
+checkpoint['cfg']['task'].bpe_tokenizer = task.build_bpe(checkpoint['cfg']['model'])
+tokenizer = checkpoint['cfg']['task'].bpe_tokenizer
+
+processor = LabelEncoder(task.dicts['text'])
+
+vocoder = SpeechT5HifiGan.from_pretrained('microsoft/speecht5_hifigan').to(device)
+
+def get_embs(emb_path):
+    spkembs = get_features_or_waveform(emb_path)
+    spkembs = torch.from_numpy(spkembs).float().unsqueeze(0)
+    return spkembs
+
+def process_text(text):
+    text = araby.strip_diacritics(text)
+    return processor(tokenizer.encode(text)).reshape(1, -1)
+
+net_input = {}
+
+def inference(text, spkr=emb_path):
+    net_input['src_tokens'] = process_text(text)
+    net_input['spkembs'] = get_embs(spkr)
+    outs, _, attn = task.generate_speech(
+            [model], 
+            net_input,
+        )
+    with torch.no_grad():
+        gen_audio = vocoder(outs.to(device))
+    return (16000,gen_audio.cpu().numpy())
+
+text_box = gr.Textbox(max_lines=2, label="Arabic Text")
+out = gr.Audio(label="Synthesized Audio", type="numpy")
+demo = gr.Interface(inference, \
+    inputs=text_box, outputs=out, title="ArTST")
+
+if __name__ == "__main__":
+    demo.launch(share=True)

artst/__init__.py

@@ -0,0 +1 @@
+from . import data, tasks, criterions, models   # noqa

artst/criterions/__init__.py

@@ -0,0 +1,10 @@
+import importlib
+import os
+
+
+for file in os.listdir(os.path.dirname(__file__)):
+    if file.endswith(".py") and not file.startswith("_"):
+        criterion_name = file[: file.find(".py")]
+        importlib.import_module(
+            "artst.criterions." + criterion_name
+        )

artst/criterions/artst_criterion.py

@@ -0,0 +1,443 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transform (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import re
+from dataclasses import dataclass
+
+import math
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from artst.criterions.text_to_speech_loss import TexttoSpeechLoss
+from artst.criterions.text_pretrain_criterion import TextPretrainCriterion, TextPretrainCriterionConfig
+from fairseq.criterions.label_smoothed_cross_entropy import LabelSmoothedCrossEntropyCriterionConfig
+from artst.criterions.speech_pretrain_criterion import SpeechPretrainCriterion, SpeechPretrainCriterionConfig
+from artst.criterions.speech_to_text_loss import SpeechtoTextLoss, SpeechtoTextLossConfig                                                  
+from fairseq.logging.meters import safe_round
+
+@dataclass
+class ArTSTCriterionConfig(
+    LabelSmoothedCrossEntropyCriterionConfig, 
+    TextPretrainCriterionConfig,
+    SpeechPretrainCriterionConfig,
+    SpeechtoTextLossConfig
+    ):
+    pass
+
+@register_criterion(
+    "artst", dataclass=ArTSTCriterionConfig
+)
+class ArTSTCriterion(FairseqCriterion):
+    def __init__(
+        self,
+        task,
+        sentence_avg,
+        label_smoothing,
+        pred_masked_weight, 
+        pred_nomask_weight, 
+        loss_weights=None, 
+        log_keys=None,
+        ignore_prefix_size=0,
+        report_accuracy=False,
+        use_masking=True,
+        use_weighted_masking=False,
+        loss_type="L1",
+        bce_pos_weight=5.0,
+        bce_loss_lambda=1.0,
+        use_guided_attn_loss=False,
+        num_heads_applied_guided_attn=2,
+        ce_weight=1.0,
+        ctc_weight=0.0,
+        hubert_weight=1.0,
+        dec_weight=1.0,
+        bart_weight=1.0,
+    ):
+        super().__init__(task)
+        self.speech_criterion = TexttoSpeechLoss(
+            task,
+            sentence_avg,
+            use_masking,
+            use_weighted_masking,
+            loss_type,
+            bce_pos_weight,
+            bce_loss_lambda,
+            use_guided_attn_loss,
+            num_heads_applied_guided_attn=num_heads_applied_guided_attn,
+        )
+        self.text_criterion = SpeechtoTextLoss(
+            SpeechtoTextLossConfig,
+            task,
+            sentence_avg,
+            label_smoothing,
+            ignore_prefix_size,
+            report_accuracy,
+            ce_weight,
+            ctc_weight
+        )
+        self.text_pretrain_criterion = TextPretrainCriterion(
+            task,
+            sentence_avg,
+            bart_weight,
+            loss_weights,
+        )
+        self.speech_pretrain_criterion = SpeechPretrainCriterion(
+            task,
+            sentence_avg,
+            pred_masked_weight,
+            pred_nomask_weight,
+            loss_weights,
+            log_keys,
+            use_masking,
+            use_weighted_masking,
+            loss_type,
+            bce_pos_weight,
+            hubert_weight,
+            dec_weight
+        )
+
+    def forward(self, model, sample, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+
+        task_name = sample['task_name']
+        if task_name == 's2t' or task_name == 's2c':
+            return self.text_criterion(model, sample, reduce)
+        elif task_name == 't2s' or task_name == 's2s':
+            return self.speech_criterion(model, sample)
+        elif task_name == 'text_pretrain':
+            return self.text_pretrain_criterion(model, sample, reduce)
+        elif task_name == 'speech_pretrain':
+            return self.speech_pretrain_criterion(model, sample, reduce)
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs):
+        """Aggregate logging outputs from data parallel training."""
+        logging_outputs_dict = {}
+        for logging_output in logging_outputs:
+            for task_name in logging_output:
+                if task_name not in ['s2t', 't2s', 's2c', 's2s', 'text_pretrain', 'speech_pretrain']:
+                    continue
+
+                if task_name not in logging_outputs_dict:
+                    logging_outputs_dict[task_name] = []
+                logging_outputs_dict[task_name].append(logging_output[task_name])
+
+        for task_name in logging_outputs_dict:
+            if task_name == 's2t':
+                # LabelSmoothedCrossEntropyCriterion.reduce_metrics([logging_output['s2t'] for logging_output in logging_outputs])
+                s2t_logging_output = logging_outputs_dict[task_name]
+                # s2t_sum = sum(log.get("ce_loss", 0) for log in logging_outputs)
+                loss_sum = sum(log.get("loss", 0) for log in s2t_logging_output)
+                nll_loss_sum = sum(log.get("nll_loss", 0) for log in s2t_logging_output)
+                ntokens = sum(log.get("ntokens", 0) for log in s2t_logging_output)
+                ce_loss_sum = sum(log.get("ce_loss", 0) for log in s2t_logging_output)
+                ctc_loss_sum = sum(log.get("ctc_loss", 0) for log in s2t_logging_output)
+
+                sample_size = max(1, sum(log.get("sample_size", 0) for log in s2t_logging_output))
+                metrics.log_scalar(
+                    "s2t_loss", loss_sum / sample_size / math.log(2), sample_size, 1, round=3
+                )
+
+                metrics.log_scalar(
+                    "s2t_nll_loss", nll_loss_sum / ntokens / math.log(2), ntokens, 2, round=3
+                )
+                metrics.log_derived(
+                    "s2t_ppl", lambda meters: utils.get_perplexity(meters["s2t_nll_loss"].avg, 2)
+                )
+                metrics.log_scalar(
+                    "ctc_loss", ctc_loss_sum / sample_size / math.log(2), ntokens, 2, round=3
+                )
+                metrics.log_scalar(
+                    "ce_loss", ce_loss_sum / ntokens, ntokens, 2, round=3
+                )
+
+                total = utils.item(sum(log.get("total", 0) for log in s2t_logging_output))
+                if total > 0:
+                    metrics.log_scalar("s2t_total", total)
+                    n_correct = utils.item(
+                        sum(log.get("n_correct", 0) for log in s2t_logging_output)
+                    )
+                    metrics.log_scalar("s2t_n_correct", n_correct)
+                    metrics.log_derived(
+                        "s2t_accuracy",
+                        lambda meters: round(
+                            meters["s2t_n_correct"].sum * 100.0 / meters["s2t_total"].sum, 3
+                        )
+                        if meters["s2t_total"].sum > 0
+                        else float("nan"),
+                        2
+                    )
+                c_errors = sum(log.get("c_errors", 0) for log in s2t_logging_output)
+                metrics.log_scalar("_c_errors", c_errors)
+                c_total = sum(log.get("c_total", 0) for log in s2t_logging_output)
+                metrics.log_scalar("_c_total", c_total)
+                w_errors = sum(log.get("w_errors", 0) for log in s2t_logging_output)
+                metrics.log_scalar("_w_errors", w_errors)
+                wv_errors = sum(log.get("wv_errors", 0) for log in s2t_logging_output)
+                metrics.log_scalar("_wv_errors", wv_errors)
+                w_total = sum(log.get("w_total", 0) for log in s2t_logging_output)
+                metrics.log_scalar("_w_total", w_total)
+                if c_total > 0:
+                    metrics.log_derived(
+                        "uer",
+                        lambda meters: safe_round(
+                            meters["_c_errors"].sum * 100.0 / meters["_c_total"].sum, 3
+                        )
+                        if meters["_c_total"].sum > 0
+                        else float("nan"),
+                    )
+                if w_total > 0:
+                    metrics.log_derived(
+                        "wer",
+                        lambda meters: safe_round(
+                            meters["_w_errors"].sum * 100.0 / meters["_w_total"].sum, 3
+                        )
+                        if meters["_w_total"].sum > 0
+                        else float("nan"),
+                    )
+                    metrics.log_derived(
+                        "raw_wer",
+                        lambda meters: safe_round(
+                            meters["_wv_errors"].sum * 100.0 / meters["_w_total"].sum, 3
+                        )
+                        if meters["_w_total"].sum > 0
+                        else float("nan"),
+                    )
+
+            if task_name == 't2s':
+                # TTSLossCriterion.reduce_metrics([logging_output['t2s'] for logging_output in logging_outputs])
+                # t2s_sum = sum(log.get("speech_loss", 0) for log in logging_outputs)
+                t2s_logging_output = logging_outputs_dict[task_name]
+                loss_sum = sum(log.get("loss", 0) for log in t2s_logging_output)
+                l1_loss_sum = sum(log.get("l1_loss", 0) for log in t2s_logging_output)
+                l2_loss_sum = sum(log.get("l2_loss", 0) for log in t2s_logging_output)
+                bce_loss_sum = sum(log.get("bce_loss", 0) for log in t2s_logging_output)
+                sample_size = max(1, sum(log.get("sample_size", 0) for log in t2s_logging_output))
+                metrics.log_scalar(
+                    "t2s_loss", loss_sum / sample_size, sample_size, 1, round=5
+                )
+                encoder_alpha_sum = sum(log.get("encoder_alpha", 0) for log in t2s_logging_output)
+                decoder_alpha_sum = sum(log.get("decoder_alpha", 0) for log in t2s_logging_output)
+                ngpu = sum(log.get("ngpu", 0) for log in t2s_logging_output)
+
+                metrics.log_scalar(
+                    "t2s_l1_loss", l1_loss_sum / sample_size, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "t2s_l2_loss", l2_loss_sum / sample_size, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "t2s_bce_loss", bce_loss_sum / sample_size, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "t2s_encoder_alpha", encoder_alpha_sum / sample_size, sample_size, round=5
+                )
+                metrics.log_scalar(
+                    "t2s_decoder_alpha", decoder_alpha_sum / sample_size, sample_size, round=5
+                )
+
+                if "enc_dec_attn_loss" in t2s_logging_output[0]:
+                    enc_dec_attn_loss_sum = sum(log.get("enc_dec_attn_loss", 0) for log in t2s_logging_output)
+                    metrics.log_scalar(
+                        "t2s_enc_dec_attn_loss", enc_dec_attn_loss_sum / sample_size, sample_size, round=8
+                    )
+
+            if task_name == 's2c':
+                s2c_logging_output = logging_outputs_dict[task_name]
+                loss_sum = sum(log.get("loss", 0) for log in s2c_logging_output)
+                nll_loss_sum = sum(log.get("nll_loss", 0) for log in s2c_logging_output)
+                ntokens = sum(log.get("ntokens", 0) for log in s2c_logging_output)
+
+                sample_size = max(1, sum(log.get("sample_size", 0) for log in s2c_logging_output))
+                metrics.log_scalar(
+                    "s2c_loss", loss_sum / sample_size / math.log(2), sample_size, 1, round=3
+                )
+
+                metrics.log_scalar(
+                    "s2c_nll_loss", nll_loss_sum / ntokens / math.log(2), ntokens, 2, round=3
+                )
+
+                total = utils.item(sum(log.get("total", 0) for log in s2c_logging_output)) 
+                if total > 0:
+                    metrics.log_scalar("s2c_total", total)
+                    n_correct = utils.item(sum(log.get("n_correct", 0) for log in s2c_logging_output))
+                    metrics.log_scalar("s2c_n_correct", n_correct)
+                    metrics.log_derived(
+                        "s2c_accuracy",
+                        lambda meters: round(
+                            meters["s2c_n_correct"].sum * 100.0 / meters["s2c_total"].sum, 3
+                        )
+                        if meters["s2c_total"].sum > 0
+                        else float("nan"),
+                        2
+                    )
+            
+            if task_name == 's2s':
+                s2s_logging_output = logging_outputs_dict[task_name]
+                loss_sum = sum(log.get("loss", 0) for log in s2s_logging_output)
+                l1_loss_sum = sum(log.get("l1_loss", 0) for log in s2s_logging_output)
+                l2_loss_sum = sum(log.get("l2_loss", 0) for log in s2s_logging_output)
+                bce_loss_sum = sum(log.get("bce_loss", 0) for log in s2s_logging_output)
+                sample_size = max(1, sum(log.get("sample_size", 0) for log in s2s_logging_output))
+                metrics.log_scalar(
+                    "s2s_loss", loss_sum / sample_size, sample_size, 1, round=5
+                )
+                encoder_alpha_sum = sum(log.get("encoder_alpha", 0) for log in s2s_logging_output)
+                decoder_alpha_sum = sum(log.get("decoder_alpha", 0) for log in s2s_logging_output)
+                ngpu = sum(log.get("ngpu", 0) for log in s2s_logging_output)
+
+                metrics.log_scalar(
+                    "s2s_l1_loss", l1_loss_sum / sample_size, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "s2s_l2_loss", l2_loss_sum / sample_size, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "s2s_bce_loss", bce_loss_sum / sample_size, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "s2s_decoder_alpha", decoder_alpha_sum / sample_size, sample_size, round=5
+                )
+
+                if "enc_dec_attn_loss" in s2s_logging_output[0]:
+                    enc_dec_attn_loss_sum = sum(log.get("enc_dec_attn_loss", 0) for log in s2s_logging_output)
+                    metrics.log_scalar(
+                        "s2s_enc_dec_attn_loss", enc_dec_attn_loss_sum / sample_size, sample_size, round=8
+                    )
+
+            if task_name == 'text_pretrain':
+                bart_logging_output = logging_outputs_dict[task_name]
+                loss_sum = sum(log.get("loss", 0) for log in bart_logging_output)
+                ntokens = sum(log.get("ntokens", 0) for log in bart_logging_output)
+                sample_size = max(1, sum(log.get("sample_size", 0) for log in bart_logging_output))
+                bart_loss_sum = sum(log.get("bart_loss", 0) for log in bart_logging_output)
+
+                # we divide by log(2) to convert the loss from base e to base 2
+                metrics.log_scalar(
+                    "text_loss", loss_sum / sample_size / math.log(2), sample_size, round=3
+                )
+                metrics.log_scalar(
+                    "bart_loss", bart_loss_sum / sample_size / math.log(2), ntokens, 2, round=3
+                )
+                if sample_size != ntokens:
+                    metrics.log_scalar(
+                        "bart_nll_loss", bart_loss_sum / ntokens / math.log(2), ntokens, round=3
+                    )
+                    metrics.log_derived(
+                        "bart_ppl", lambda meters: utils.get_perplexity(meters["bart_nll_loss"].avg)
+                    )
+                else:
+                    metrics.log_derived(
+                        "bart_ppl", lambda meters: utils.get_perplexity(meters["bart_loss"].avg)
+                    )
+                metrics.log_scalar("bart_wpb", ntokens, priority=180, round=1)
+
+                val_prob_perplexity = 0
+                val_code_perplexity = 0
+                sample_size_pp = 0
+                count_log_cp = 0
+                for log in bart_logging_output:
+                    if "loss_prob_perplexity" in log:
+                        val_prob_perplexity = val_prob_perplexity + log["loss_prob_perplexity"]
+                        sample_size_pp = sample_size_pp + log["sample_size"]
+                    if "code_perplexity" in log:
+                        val_code_perplexity = val_code_perplexity + log["code_perplexity"]
+                        count_log_cp = count_log_cp + 1
+                if val_prob_perplexity > 0:
+                    metrics.log_scalar("text_loss_prob_perplexity", val_prob_perplexity / sample_size_pp / math.log(2), round=3)
+                if val_code_perplexity > 0:
+                    metrics.log_scalar("text_code_perplexity", val_code_perplexity / count_log_cp, round=3)
+
+            if task_name == 'speech_pretrain':
+                hubert_logging_output = logging_outputs_dict[task_name]
+                loss_sum = sum(log.get("loss", 0) for log in hubert_logging_output)
+                ntokens = sum(log.get("ntokens", 0) for log in hubert_logging_output)
+                sample_size = max(1, sum(log.get("sample_size", 0) for log in hubert_logging_output))
+                dec_loss_sum = sum(log.get("dec_loss", 0) for log in hubert_logging_output)
+                l1_loss_sum = sum(log.get("l1_loss", 0) for log in hubert_logging_output)
+                l2_loss_sum = sum(log.get("l2_loss", 0) for log in hubert_logging_output)
+                bce_loss_sum = sum(log.get("bce_loss", 0) for log in hubert_logging_output)
+                ngpu = sum(log.get("ngpu", 0) for log in hubert_logging_output)
+
+                metrics.log_scalar("hubert_loss", loss_sum / sample_size / math.log(2), sample_size, round=3)
+                if sample_size != ntokens:
+                    metrics.log_scalar("hubert_nll_loss", loss_sum / ntokens / math.log(2), ntokens, round=3)
+                    metrics.log_derived("hubert_ppl", lambda meters: utils.get_perplexity(meters["hubert_nll_loss"].avg))
+                else:
+                    metrics.log_derived("hubert_ppl", lambda meters: utils.get_perplexity(meters["hubert_loss"].avg))
+
+                counts = {}
+                for lk in hubert_logging_output[0].keys():
+                    if lk.startswith("count_"):
+                        val = sum(log[lk] for log in hubert_logging_output)
+                        metrics.log_scalar("hubert_" + lk, val)
+                        counts[lk] = val
+
+                for lk in hubert_logging_output[0].keys():
+                    if lk.startswith("loss_") and lk != 'loss_prob_perplexity':
+                        val = sum(log[lk] for log in hubert_logging_output)
+                        metrics.log_scalar("hubert_" + lk, val / sample_size / math.log(2), round=3)
+                    elif lk.startswith("correct_"):
+                        val = sum(log[lk] for log in hubert_logging_output)
+                        metrics.log_scalar("hubert_" + lk, val / counts[re.sub("correct", "count", lk)])
+                    # elif lk == 'code_perplexity':
+                    #     val = sum(log[lk] for log in hubert_logging_output)
+                    #     metrics.log_scalar("hubert_" + lk, val / len(hubert_logging_output), round=3)
+
+                val_prob_perplexity = 0
+                val_code_perplexity = 0
+                sample_size_pp = 0
+                count_log_cp = 0
+                for log in hubert_logging_output:
+                    if "loss_prob_perplexity" in log:
+                        val_prob_perplexity = val_prob_perplexity + log["loss_prob_perplexity"]
+                        sample_size_pp = sample_size_pp + log["sample_size"]
+                    if "code_perplexity" in log:
+                        val_code_perplexity = val_code_perplexity + log["code_perplexity"]
+                        count_log_cp = count_log_cp + 1
+                if val_prob_perplexity > 0:
+                    metrics.log_scalar("hubert_loss_prob_perplexity", val_prob_perplexity / sample_size_pp / math.log(2), round=3)
+                if val_code_perplexity > 0:
+                    metrics.log_scalar("hubert_code_perplexity", val_code_perplexity / count_log_cp, round=3)
+
+                metrics.log_scalar(
+                    "hubert_dec_loss", dec_loss_sum / ngpu, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "hubert_l1_loss", l1_loss_sum / ngpu, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "hubert_l2_loss", l2_loss_sum / ngpu, sample_size, 2, round=5
+                )
+                metrics.log_scalar(
+                    "hubert_bce_loss", bce_loss_sum / ngpu, sample_size, 2, round=5
+                )
+                if "enc_dec_attn_loss" in hubert_logging_output[0]:
+                    enc_dec_attn_loss_sum = sum(log.get("enc_dec_attn_loss", 0) for log in hubert_logging_output)
+                    metrics.log_scalar(
+                        "hubert_enc_dec_attn_loss", enc_dec_attn_loss_sum / ngpu, sample_size, round=8
+                    )
+                metrics.log_scalar("hubert_wpb", ntokens, priority=180, round=1)
+
+        loss = sum(log.get("loss", 0) for log in logging_outputs)
+        sample_size = max(1, sum(log.get("sample_size", 0) for log in logging_outputs))
+        metrics.log_scalar(
+            "loss", loss / sample_size, sample_size, 1, round=5
+        )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return False

artst/criterions/speech_pretrain_criterion.py

@@ -0,0 +1,265 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transform (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import math
+import re
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion
+from artst.criterions.text_to_speech_loss import TexttoSpeechLoss, TexttoSpeechLossConfig
+
+
+@dataclass
+class SpeechPretrainCriterionConfig(TexttoSpeechLossConfig):
+    pred_masked_weight: float = field(
+        default=1.0,
+        metadata={"help": "weight for predictive loss for masked frames"},
+    )
+    pred_nomask_weight: float = field(
+        default=0.0,
+        metadata={"help": "weight for predictive loss for unmasked frames"},
+    )
+    loss_weights: Optional[List[float]] = field(
+        default_factory=lambda: [10,],
+        metadata={"help": "weights for additional loss terms (not first one)"},
+    )
+    log_keys: List[str] = field(
+        default_factory=lambda: [],
+        metadata={"help": "output keys to log"},
+    )
+    hubert_weight: float = field(
+        default=1.0,
+        metadata={"help": "weight of hubert loss"},
+    )
+    dec_weight: float = field(
+        default=1.0,
+        metadata={"help": "weight of decoder loss"},
+    )
+
+
+class SpeechPretrainCriterion(FairseqCriterion):
+    def __init__(
+        self, 
+        task, 
+        sentence_avg, 
+        pred_masked_weight, 
+        pred_nomask_weight, 
+        loss_weights=None, 
+        log_keys=None,
+        use_masking=True,
+        use_weighted_masking=False,
+        loss_type="L1",
+        bce_pos_weight=5.0,
+        hubert_weight=1.0,
+        dec_weight=1.0,
+    ):
+        super().__init__(task)
+        self.pred_masked_weight = pred_masked_weight
+        self.pred_nomask_weight = pred_nomask_weight
+        self.loss_weights = loss_weights
+        self.log_keys = [] if log_keys is None else log_keys
+        self.hubert_weight = hubert_weight
+        self.dec_weight = dec_weight
+
+        self.speech_criterion = TexttoSpeechLoss(
+            task,
+            sentence_avg,
+            use_masking,
+            use_weighted_masking,
+            loss_type,
+            bce_pos_weight,
+        )
+
+    def forward(self, model, sample, reduce=True, log_pred=False):
+        """Compute the loss for the given sample.
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        if self.dec_weight == 0:
+            sample["net_input"]["only_hubert"] = True
+        net_output, net_output_dec = model(target_list=sample["target_list"], **sample["net_input"])
+        loss = 0.
+        sample_size = 0
+        logging_output = {}
+        reduction = "sum" if reduce else "none"
+
+        loss_m_list = []
+        logp_m_list = model.get_logits(net_output, True)
+        targ_m_list = model.get_targets(None, net_output, True)
+        assert self.pred_masked_weight == 0 or len(logp_m_list) > 0
+        for i, (logp_m, targ_m) in enumerate(zip(logp_m_list, targ_m_list)):
+            loss_m = F.cross_entropy(logp_m, targ_m, reduction=reduction)
+            loss_m_list.append(loss_m)
+            logging_output[f"loss_m_{i}"] = loss_m.detach().item()
+        if self.pred_masked_weight > 0:
+            loss += self.pred_masked_weight * sum(loss_m_list)
+            sample_size += targ_m_list[0].numel()
+
+        loss_u_list = []
+        logp_u_list = model.get_logits(net_output, False)
+        targ_u_list = model.get_targets(None, net_output, False)
+        assert self.pred_nomask_weight == 0 or len(logp_u_list) > 0
+        for i, (logp_u, targ_u) in enumerate(zip(logp_u_list, targ_u_list)):
+            loss_u = F.cross_entropy(logp_u, targ_u, reduction=reduction)
+            loss_u_list.append(loss_u)
+            logging_output[f"loss_u_{i}"] = loss_u.detach().item()
+        if self.pred_nomask_weight > 0:
+            loss += self.pred_nomask_weight * sum(loss_u_list)
+            sample_size += targ_u_list[0].numel()
+
+        if self.loss_weights is not None:
+            assert hasattr(model, "get_extra_losses")
+            extra_losses, names = model.get_extra_losses(net_output)
+            if torch.is_tensor(extra_losses):
+                extra_losses = [extra_losses]
+                names = [names]
+            if len(self.loss_weights) == 1 and len(extra_losses) != 1:
+                self.loss_weights = [self.loss_weights[0]] * len(extra_losses)
+            if len(self.loss_weights) > len(extra_losses):
+                modified_loss_weight = self.loss_weights[:len(extra_losses)]
+            else:
+                modified_loss_weight = self.loss_weights
+
+            # assert len(extra_losses) == len(self.loss_weights), f"{len(extra_losses)}, {len(self.loss_weights)}"
+            for p, n, coef in zip(extra_losses, names, modified_loss_weight):
+                # print(n + str(coef))
+                if coef != 0 and p is not None:
+                    p = coef * p.float() * sample_size
+                    loss += p
+                    logging_output[f"loss_{n}"] = p.detach().item()
+
+        logging_output = {
+            "ntokens": sample_size,
+            "nsentences": sample["id"].numel(),
+            "sample_size": sample_size,
+            "ngpu": 1,
+            **logging_output,
+        }
+
+        if 'loss_prob_perplexity' in logging_output:
+            logging_output['code_perplexity'] = net_output['code_perplexity'].detach().item()     
+
+        for lk in self.log_keys:
+            if lk in net_output:
+                logging_output[lk] = float((net_output[lk].item()))
+
+        def compute_correct(logits):
+            if logits.numel() == 0:
+                return 0, 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 = max.numel()
+                return corr, count
+
+        with torch.no_grad():
+            for i, logp_m in enumerate(logp_m_list):
+                corr_m, count_m = compute_correct(logp_m)
+                logging_output[f"correct_m_{i}"] = corr_m
+                logging_output[f"count_m_{i}"] = count_m
+
+            for i, logp_u in enumerate(logp_u_list):
+                corr_u, count_u = compute_correct(logp_u)
+                logging_output[f"correct_u_{i}"] = corr_u
+                logging_output[f"count_u_{i}"] = count_u
+
+        if self.dec_weight == 0.0:
+            logging_output["loss"] = loss.item() if reduce else loss
+            return loss, sample_size, logging_output
+
+#       ## dec loss
+        dec_loss, l1_loss, l2_loss, bce_loss, enc_dec_attn_loss = self.speech_criterion.compute_loss(model, net_output_dec, sample)
+        
+        # Log tts loss
+        logging_output['dec_loss'] = dec_loss.item()
+        logging_output['l1_loss'] = l1_loss.item()
+        logging_output['l2_loss'] = l2_loss.item()
+        logging_output['bce_loss'] = bce_loss.item()
+        if enc_dec_attn_loss is not None:
+            logging_output['enc_dec_attn_loss'] = enc_dec_attn_loss.item()
+
+        loss = self.hubert_weight * loss + self.dec_weight * sample_size * dec_loss
+        logging_output["loss"] = loss.item() if reduce else loss
+        return loss, sample_size, logging_output
+
+    @staticmethod
+    def reduce_metrics(logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training (copied from normal cross entropy)."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+        dec_loss_sum = sum(log.get("dec_loss", 0) for log in logging_outputs)
+        l1_loss_sum = sum(log.get("l1_loss", 0) for log in logging_outputs)
+        l2_loss_sum = sum(log.get("l2_loss", 0) for log in logging_outputs)
+        bce_loss_sum = sum(log.get("bce_loss", 0) for log in logging_outputs)
+        ngpu = sum(log.get("ngpu", 0) for log in logging_outputs)
+
+        metrics.log_scalar("loss", loss_sum / sample_size / math.log(2), sample_size, round=3)
+        if sample_size != ntokens:
+            metrics.log_scalar("nll_loss", loss_sum / ntokens / math.log(2), ntokens, round=3)
+            metrics.log_derived("ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg))
+        else:
+            metrics.log_derived("ppl", lambda meters: utils.get_perplexity(meters["loss"].avg))
+
+        counts = {}
+        for lk in logging_outputs[0].keys():
+            if lk.startswith("count_"):
+                val = sum(log[lk] for log in logging_outputs)
+                metrics.log_scalar(lk, val)
+                counts[lk] = val
+
+        for lk in logging_outputs[0].keys():
+            if lk.startswith("loss_"):
+                val = sum(log[lk] for log in logging_outputs)
+                metrics.log_scalar(lk, val / sample_size / math.log(2), round=3)
+            elif lk.startswith("correct_"):
+                val = sum(log[lk] for log in logging_outputs)
+                metrics.log_scalar(lk, val / counts[re.sub("correct", "count", lk)])
+            elif lk == 'code_perplexity':
+                val = sum(log[lk] for log in logging_outputs)
+                metrics.log_scalar(lk, val / len(logging_outputs), round=3)
+
+        metrics.log_scalar(
+            "dec_loss", dec_loss_sum / ngpu, sample_size, 2, round=5
+        )
+        metrics.log_scalar(
+            "l1_loss", l1_loss_sum / ngpu, sample_size, 2, round=5
+        )
+        metrics.log_scalar(
+            "l2_loss", l2_loss_sum / ngpu, sample_size, 2, round=5
+        )
+        metrics.log_scalar(
+            "bce_loss", bce_loss_sum / ngpu, sample_size, 2, round=5
+        )
+        if "enc_dec_attn_loss" in logging_outputs[0]:
+            enc_dec_attn_loss_sum = sum(log.get("enc_dec_attn_loss", 0) for log in logging_outputs)
+            metrics.log_scalar(
+                "enc_dec_attn_loss", enc_dec_attn_loss_sum / ngpu, sample_size, round=8
+            )
+
+    @staticmethod
+    def aggregate_logging_outputs(logging_outputs):
+        """Aggregate logging outputs from data parallel training."""
+        raise NotImplementedError()
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return False

artst/criterions/speech_to_text_loss.py

@@ -0,0 +1,473 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transform (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import math
+from argparse import Namespace
+from dataclasses import dataclass, field
+from omegaconf import II
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from fairseq.data.data_utils import post_process
+from fairseq.tasks import FairseqTask
+from fairseq.logging.meters import safe_round
+
+import logging
+logger = logging.getLogger(__name__)
+
+@dataclass
+class SpeechtoTextLossConfig(FairseqDataclass):
+    zero_infinity: bool = field(
+        default=False,
+        metadata={"help": "zero inf loss when source length <= target length"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+    post_process: Optional[str] = field(
+        default="sentencepiece",
+        metadata={
+            "help": "how to post process predictions into words. can be letter, "
+            "wordpiece, BPE symbols, etc. "
+            "See fairseq.data.data_utils.post_process() for full list of options"
+        },
+    )
+    wer_kenlm_model: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "if this is provided, use kenlm to compute wer (along with other wer_* args)"
+        },
+    )
+    wer_lexicon: Optional[str] = field(
+        default=None,
+        metadata={"help": "lexicon to use with wer_kenlm_model"},
+    )
+    wer_lm_weight: float = field(
+        default=2.0,
+        metadata={"help": "lm weight to use with wer_kenlm_model"},
+    )
+    wer_word_score: float = field(
+        default=-1.0,
+        metadata={"help": "lm word score to use with wer_kenlm_model"},
+    )
+
+    wer_args: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "DEPRECATED: tuple of (wer_kenlm_model, wer_lexicon, wer_lm_weight, wer_word_score)"
+        },
+    )
+
+    label_smoothing: float = field(
+        default=0.0,
+        metadata={"help": "epsilon for label smoothing, 0 means no label smoothing"},
+    )
+    report_accuracy: bool = field(
+        default=False,
+        metadata={"help": "report accuracy metric"},
+    )
+    ignore_prefix_size: int = field(
+        default=0,
+        metadata={"help": "Ignore first N tokens"},
+    )
+    #: bool = II("optimization.sentence_avg")
+
+    ce_weight: float = field(
+        default=1.0,
+        metadata={"help": "loss weight for cross entropy"},
+    )
+    ctc_weight: float = field(
+        default=0.0,
+        metadata={"help": "loss weiehgt for ctc in ASR"},
+    )
+
+
+def label_smoothed_nll_loss(lprobs, target, epsilon, ignore_index=None, reduce=True):
+    if target.dim() == lprobs.dim() - 1:
+        target = target.unsqueeze(-1)
+    nll_loss = -lprobs.gather(dim=-1, index=target)
+    smooth_loss = -lprobs.sum(dim=-1, keepdim=True)
+    if ignore_index is not None:
+        pad_mask = target.eq(ignore_index)
+        nll_loss.masked_fill_(pad_mask, 0.0)
+        smooth_loss.masked_fill_(pad_mask, 0.0)
+    else:
+        nll_loss = nll_loss.squeeze(-1)
+        smooth_loss = smooth_loss.squeeze(-1)
+    if reduce:
+        nll_loss = nll_loss.sum()
+        smooth_loss = smooth_loss.sum()
+    eps_i = epsilon / (lprobs.size(-1) - 1)
+    loss = (1.0 - epsilon - eps_i) * nll_loss + eps_i * smooth_loss
+    return loss, nll_loss
+
+
+class SpeechtoTextLoss(FairseqCriterion):
+    def __init__(
+        self,
+        cfg: SpeechtoTextLossConfig, 
+        task: FairseqTask,
+        sentence_avg=True,
+        label_smoothing=0.1,
+        ignore_prefix_size=0,
+        report_accuracy=False,
+        ce_weight=1.0,
+        ctc_weight=0.0,
+    ):
+
+        super().__init__(task)
+        self.blank_idx = (
+            task.target_dictionary.index(task.blank_symbol)
+            if hasattr(task, "blank_symbol")
+            else 0
+        )
+        #print ("self.blank_idx: ", self.blank_idx)
+
+        self.pad_idx = task.target_dictionary.pad()
+        self.eos_idx = task.target_dictionary.eos()
+        self.post_process = cfg.post_process
+        self.ce_weight = ce_weight
+        self.ctc_weight = ctc_weight
+
+        ## for ce
+        self.sentence_avg = sentence_avg
+        self.eps = label_smoothing
+        self.ignore_prefix_size = ignore_prefix_size
+        self.report_accuracy = report_accuracy
+
+        if cfg.wer_args is not None:
+            (
+                cfg.wer_kenlm_model,
+                cfg.wer_lexicon,
+                cfg.wer_lm_weight,
+                cfg.wer_word_score,
+            ) = eval(cfg.wer_args)
+
+        if cfg.wer_kenlm_model is not None:
+            from examples.speech_recognition.w2l_decoder import W2lKenLMDecoder
+
+            dec_args = Namespace()
+            dec_args.nbest = 1
+            dec_args.criterion = "ctc"
+            dec_args.kenlm_model = cfg.wer_kenlm_model
+            dec_args.lexicon = cfg.wer_lexicon
+            dec_args.beam = 50
+            dec_args.beam_size_token = min(50, len(task.target_dictionary))
+            dec_args.beam_threshold = min(50, len(task.target_dictionary))
+            dec_args.lm_weight = cfg.wer_lm_weight
+            dec_args.word_score = cfg.wer_word_score
+            dec_args.unk_weight = -math.inf
+            dec_args.sil_weight = 0
+
+            self.w2l_decoder = W2lKenLMDecoder(dec_args, task.target_dictionary)
+        else:
+            self.w2l_decoder = None
+
+        self.zero_infinity = cfg.zero_infinity
+        #self.sentence_avg = cfg.sentence_avg
+
+        if self.ce_weight > 0 and self.ctc_weight > 0:
+            logger.info("Using cross entropy loss and CTC loss for ASR")
+        elif self.ce_weight > 0:
+            logger.info("Only using CE loss")
+        elif self.ctc_weight > 0:
+            logger.info("Only using CTC loss for ASR")
+        else:
+            logger.info("ERROR")
+
+    def forward(self, model, sample, reduce=True):
+
+        if self.ce_weight == 0 and self.ctc_weight > 0:
+            sample["only_ctc"] = True
+
+        net_output_decoder, net_output = model(**sample["net_input"])
+        
+        if self.ce_weight > 0:
+            loss_ce, nll_loss_ce = self.compute_loss(model, net_output_decoder, sample, reduce=reduce)
+            #print ("loss_ce: ", loss_ce)
+        else:
+            nll_loss_ce = None
+
+        if self.ctc_weight > 0:
+            loss_ctc, lprobs, input_lengths = self.compute_loss_ctc(model, net_output, sample)
+
+        if self.ce_weight > 0 and self.ctc_weight > 0:
+            loss = self.ce_weight * loss_ce + self.ctc_weight * loss_ctc
+        elif self.ce_weight > 0:
+            loss = loss_ce
+        elif self.ctc_weight > 0:
+            loss = loss_ctc
+        else:
+            logger.info("ERROR: must ce_weight > 0 or ctc_weight > 0")
+
+        ntokens = (
+            sample["ntokens"] if "ntokens" in sample else sample["target_lengths"].sum().item()
+        )
+
+        sample_size = sample["target"].size(0) if self.sentence_avg else ntokens
+
+        logging_output = {
+            "loss": loss.item(),
+            "ce_loss": loss_ce.item() if self.ce_weight > 0 else 0,
+            "ctc_loss": loss_ctc.item() if self.ctc_weight > 0 else 0,
+            "nll_loss": nll_loss_ce.item() if nll_loss_ce is not None else 0,
+            "ntokens": sample["ntokens"],
+            "nsentences": sample["target"].size(0),
+            "sample_size": sample_size,
+        }
+
+        if self.ce_weight > 0 and self.report_accuracy:
+            n_correct, total = self.compute_accuracy(model, net_output_decoder, sample)
+            logging_output["n_correct"] = utils.item(n_correct.item())
+            logging_output["total"] = utils.item(total.data)
+
+        if self.ctc_weight > 0 and not model.training:
+            import editdistance
+
+            with torch.no_grad():
+                lprobs_t = lprobs.transpose(0, 1).float().contiguous().cpu()
+
+                c_err = 0
+                c_len = 0
+                w_errs = 0
+                w_len = 0
+                wv_errs = 0
+                for lp, t, inp_l in zip(
+                    lprobs_t,
+                    sample["target_label"]
+                    if "target_label" in sample
+                    else sample["target"],
+                    input_lengths,
+                ):
+                    lp = lp[:inp_l].unsqueeze(0)
+
+                    decoded = None
+                    if self.w2l_decoder is not None:
+                        decoded = self.w2l_decoder.decode(lp)
+                        if len(decoded) < 1:
+                            decoded = None
+                        else:
+                            decoded = decoded[0]
+                            if len(decoded) < 1:
+                                decoded = None
+                            else:
+                                decoded = decoded[0]
+
+                    p = (t != self.task.target_dictionary.pad()) & (
+                        t != self.task.target_dictionary.eos()
+                    )
+                    targ = t[p]
+                    targ_units = self.task.target_dictionary.string(targ)
+                    targ_units_arr = targ.tolist()
+
+                    toks = lp.argmax(dim=-1).unique_consecutive()
+                    pred_units_arr = toks[toks != self.blank_idx].tolist()
+
+                    c_err += editdistance.eval(pred_units_arr, targ_units_arr)
+                    c_len += len(targ_units_arr)
+
+                    targ_words = post_process(targ_units, self.post_process).split()
+
+                    pred_units = self.task.target_dictionary.string(pred_units_arr)
+                    pred_words_raw = post_process(pred_units, self.post_process).split()
+
+                    if decoded is not None and "words" in decoded:
+                        pred_words = decoded["words"]
+                        w_errs += editdistance.eval(pred_words, targ_words)
+                        wv_errs += editdistance.eval(pred_words_raw, targ_words)
+                    else:
+                        dist = editdistance.eval(pred_words_raw, targ_words)
+                        w_errs += dist
+                        wv_errs += dist
+
+                    w_len += len(targ_words)
+
+                logging_output["wv_errors"] = wv_errs
+                logging_output["w_errors"] = w_errs
+                logging_output["w_total"] = w_len
+                logging_output["c_errors"] = c_err
+                logging_output["c_total"] = c_len
+
+        return loss, sample_size, logging_output
+
+    def compute_loss_ctc(self, model, net_output, sample):
+        lprobs = model.get_normalized_probs_for_ctc(
+            net_output, log_probs=True
+        ).contiguous()  # (T, B, C) from the encoder
+
+        if net_output["encoder_padding_mask"] is not None:
+            non_padding_mask = ~net_output["encoder_padding_mask"][0]
+            input_lengths = non_padding_mask.long().sum(-1)
+        else:
+            input_lengths = lprobs.new_full(
+                (lprobs.size(1),), lprobs.size(0), dtype=torch.long
+            )
+
+        pad_mask = (sample["target"] != self.pad_idx) & (
+            sample["target"] != self.eos_idx
+        )
+        targets_flat = sample["target"].masked_select(pad_mask)
+        if "target_lengths" in sample:
+            target_lengths = sample["target_lengths"]
+        else:
+            target_lengths = pad_mask.sum(-1)
+
+        ##processing
+        target_lengths = target_lengths - 1
+
+        with torch.backends.cudnn.flags(enabled=False):
+            loss_ctc = F.ctc_loss(
+                lprobs,
+                targets_flat,
+                input_lengths,
+                target_lengths,
+                blank=self.blank_idx,
+                reduction="sum",
+                zero_infinity=True,
+            )
+
+        return loss_ctc, lprobs, input_lengths
+
+    ## for ce
+    def get_lprobs_and_target(self, model, net_output, sample):
+        lprobs = model.get_normalized_probs(net_output, log_probs=True)
+        target = model.get_targets(sample, net_output)
+        if self.ignore_prefix_size > 0:
+            if getattr(lprobs, "batch_first", False):
+                lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
+                target = target[:, self.ignore_prefix_size :].contiguous()
+            else:
+                lprobs = lprobs[self.ignore_prefix_size :, :, :].contiguous()
+                target = target[self.ignore_prefix_size :, :].contiguous()
+        return lprobs.view(-1, lprobs.size(-1)), target.view(-1)
+
+    def compute_loss(self, model, net_output, sample, reduce=True):
+        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
+        loss, nll_loss = label_smoothed_nll_loss(
+            lprobs,
+            target,
+            self.eps,
+            ignore_index=self.padding_idx,
+            reduce=reduce,
+        )
+        return loss, nll_loss
+
+    def compute_accuracy(self, model, net_output, sample):
+        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
+        mask = target.ne(self.padding_idx)
+        n_correct = torch.sum(
+            lprobs.argmax(1).masked_select(mask).eq(target.masked_select(mask))
+        )
+        total = torch.sum(mask)
+        return n_correct, total
+
+
+    @staticmethod
+    def reduce_metrics(logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+
+        loss_sum = utils.item(sum(log.get("loss", 0) for log in logging_outputs))
+        nll_loss_sum = sum(log.get("nll_loss", 0) for log in logging_outputs)
+        ce_loss_sum = sum(log.get("ce_loss", 0) for log in logging_outputs)
+        ctc_loss_sum = sum(log.get("ctc_loss", 0) for log in logging_outputs)
+        ntokens = utils.item(sum(log.get("ntokens", 0) for log in logging_outputs))
+        nsentences = utils.item(
+            sum(log.get("nsentences", 0) for log in logging_outputs)
+        )
+        sample_size = utils.item(
+            sum(log.get("sample_size", 0) for log in logging_outputs)
+        )
+
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size / math.log(2), sample_size, round=3
+        )
+
+        metrics.log_scalar(
+            "ctc_loss", ctc_loss_sum / sample_size / math.log(2), ntokens, 2, round=3
+        )
+        metrics.log_scalar(
+            "ce_loss", ce_loss_sum / ntokens, ntokens, 2, round=3
+        )
+        metrics.log_scalar(
+            "nll_loss", nll_loss_sum / ntokens / math.log(2), ntokens, 2, round=3
+        )
+        metrics.log_derived(
+            "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg, 2)
+        )
+        
+        total = utils.item(sum(log.get("total", 0) for log in logging_outputs))
+        if total > 0:
+            metrics.log_scalar("total", total)
+            n_correct = utils.item(
+                sum(log.get("n_correct", 0) for log in logging_outputs)
+            )
+            metrics.log_scalar("n_correct", n_correct)
+            metrics.log_derived(
+                "accuracy",
+                lambda meters: round(
+                    meters["n_correct"].sum * 100.0 / meters["total"].sum, 3
+                )
+                if meters["total"].sum > 0
+                else float("nan"),
+                2
+            )
+
+        metrics.log_scalar("ntokens", ntokens)
+        metrics.log_scalar("nsentences", nsentences)
+        if sample_size != ntokens:
+            metrics.log_scalar(
+                "nll_loss", loss_sum / ntokens / math.log(2), ntokens, round=3
+            )
+
+        c_errors = sum(log.get("c_errors", 0) for log in logging_outputs)
+        metrics.log_scalar("_c_errors", c_errors)
+        c_total = sum(log.get("c_total", 0) for log in logging_outputs)
+        metrics.log_scalar("_c_total", c_total)
+        w_errors = sum(log.get("w_errors", 0) for log in logging_outputs)
+        metrics.log_scalar("_w_errors", w_errors)
+        wv_errors = sum(log.get("wv_errors", 0) for log in logging_outputs)
+        metrics.log_scalar("_wv_errors", wv_errors)
+        w_total = sum(log.get("w_total", 0) for log in logging_outputs)
+        metrics.log_scalar("_w_total", w_total)
+
+        if c_total > 0:
+            metrics.log_derived(
+                "uer",
+                lambda meters: safe_round(
+                    meters["_c_errors"].sum * 100.0 / meters["_c_total"].sum, 3
+                )
+                if meters["_c_total"].sum > 0
+                else float("nan"),
+            )
+        if w_total > 0:
+            metrics.log_derived(
+                "wer",
+                lambda meters: safe_round(
+                    meters["_w_errors"].sum * 100.0 / meters["_w_total"].sum, 3
+                )
+                if meters["_w_total"].sum > 0
+                else float("nan"),
+            )
+            metrics.log_derived(
+                "raw_wer",
+                lambda meters: safe_round(
+                    meters["_wv_errors"].sum * 100.0 / meters["_w_total"].sum, 3
+                )
+                if meters["_w_total"].sum > 0
+                else float("nan"),
+            )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

artst/criterions/text_pretrain_criterion.py

@@ -0,0 +1,142 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transform (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import math
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from omegaconf import II
+
+
+@dataclass
+class TextPretrainCriterionConfig(FairseqDataclass):
+    sentence_avg: bool = II("optimization.sentence_avg")
+    loss_weights: Optional[List[float]] = field(
+        default_factory=lambda: [0.1,],
+        metadata={"help": "weights for additional loss terms (not first one)"},
+    )
+    bart_weight: float = field(
+        default=1.0,
+        metadata={"help": "loss weight for cross entropy"},
+    )
+
+
+class TextPretrainCriterion(FairseqCriterion):
+    def __init__(self, task, sentence_avg, bart_weight, loss_weights=None):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.loss_weights = loss_weights
+        self.bart_weight = bart_weight
+
+    def forward(self, model, sample, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        net_output, codebook_out, encoder_output = model(**sample["net_input"])
+        bart_loss, _ = self.compute_loss(model, net_output, sample, reduce=reduce)
+        sample_size = (
+            sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
+        )
+        
+        loss = self.bart_weight * bart_loss
+        logging_output = {
+            "loss": loss.item(),
+            "ntokens": sample["ntokens"],
+            "nsentences": sample["target"].size(0),
+            "bart_loss": bart_loss.item(),
+            "sample_size": sample_size,
+        }
+
+        if "prob_perplexity" in codebook_out:
+            assert hasattr(model, "get_extra_losses")
+            extra_losses, names = model.get_extra_losses(codebook_out)
+            if torch.is_tensor(extra_losses):
+                extra_losses = [extra_losses]
+                names = [names]
+            if len(self.loss_weights) == 1 and len(extra_losses) != 1:
+                self.loss_weights = [self.loss_weights[0]] * len(extra_losses)
+            if len(self.loss_weights) > len(extra_losses):
+                modified_loss_weight = self.loss_weights[len(extra_losses):]
+            else:
+                modified_loss_weight = self.loss_weights
+
+            # assert len(extra_losses) == len(self.loss_weights), f"{len(extra_losses)}, {len(self.loss_weights)}"
+            for p, n, coef in zip(extra_losses, names, modified_loss_weight):
+                # print(n + str(coef))
+                if coef != 0 and p is not None:
+                    p = coef * p.float() * sample_size
+                    loss += p
+                    logging_output[f"loss_{n}"] = p.item()
+
+        if 'loss_prob_perplexity' in logging_output:
+            logging_output['code_perplexity'] = codebook_out['code_perplexity'].item()
+
+        return loss, sample_size, logging_output
+
+    def compute_loss(self, model, net_output, sample, reduce=True):
+        lprobs = model.get_normalized_probs(net_output, log_probs=True)
+        lprobs = lprobs.view(-1, lprobs.size(-1))
+        target = model.get_targets(sample, net_output).view(-1)
+        loss = F.nll_loss(
+            lprobs,
+            target,
+            ignore_index=self.padding_idx,
+            reduction="sum" if reduce else "none",
+        )
+        return loss, loss
+
+    @staticmethod
+    def reduce_metrics(logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+        bart_loss_sum = sum(log.get("bart_loss", 0) for log in logging_outputs)
+
+        # we divide by log(2) to convert the loss from base e to base 2
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size / math.log(2), sample_size, round=3
+        )
+        metrics.log_scalar(
+            "bart_loss", bart_loss_sum / sample_size / math.log(2), ntokens, 2, round=3
+        )
+        if sample_size != ntokens:
+            metrics.log_scalar(
+                "nll_loss", bart_loss_sum / ntokens / math.log(2), ntokens, round=3
+            )
+            metrics.log_derived(
+                "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
+            )
+        else:
+            metrics.log_derived(
+                "ppl", lambda meters: utils.get_perplexity(meters["bart_loss"].avg)
+            )
+
+        if "loss_prob_perplexity" in logging_outputs[0].keys():
+            val = sum(log["loss_prob_perplexity"] for log in logging_outputs)
+            metrics.log_scalar("loss_prob_perplexity", val / sample_size / math.log(2), round=3)
+        if "code_perplexity" in logging_outputs[0].keys():
+            val = sum(log["code_perplexity"] for log in logging_outputs)
+            metrics.log_scalar("code_perplexity", val / len(logging_outputs), round=3)
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

artst/criterions/text_to_speech_loss.py

@@ -0,0 +1,425 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transform (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+from dataclasses import dataclass, field
+
+import torch
+from fairseq import metrics, utils
+from espnet.nets.pytorch_backend.nets_utils import make_non_pad_mask
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from artst.models.modules.speech_encoder_prenet import SpeechEncoderPrenet
+from espnet.nets.pytorch_backend.e2e_tts_tacotron2 import GuidedAttentionLoss
+from omegaconf import II
+from typing import Any
+
+
+@dataclass
+class TexttoSpeechLossConfig(FairseqDataclass):
+    use_masking: bool = field(
+        default=True,
+        metadata={"help": "Whether to use masking in calculation of loss"},
+    )
+    use_weighted_masking: bool = field(
+        default=False,
+        metadata={"help": "Whether to use weighted masking in calculation of loss"},
+    )
+    loss_type: str = field(
+        default="L1",
+        metadata={"help": "How to calc loss"},
+    )
+    bce_pos_weight: float = field(
+        default=5.0,
+        metadata={"help": "Positive sample weight in BCE calculation (only for use-masking=True)"},
+    )
+    bce_loss_lambda: float = field(
+        default=1.0,
+        metadata={"help": "Lambda in bce loss"},
+    )
+    use_guided_attn_loss: bool = field(
+        default=False,
+        metadata={"help": "Whether to use guided attention loss"},
+    )
+    guided_attn_loss_sigma: float = field(
+        default=0.4,
+        metadata={"help": "Sigma in guided attention loss"},
+    )
+    guided_attn_loss_lambda: float = field(
+        default=10.0,
+        metadata={"help": "Lambda in guided attention loss"},
+    )
+    num_layers_applied_guided_attn: int = field(
+        default=2,
+        metadata={"help": "Number of layers to be applied guided attention loss, if set -1, all of the layers will be applied."},
+    )
+    num_heads_applied_guided_attn: int = field(
+        default=2,
+        metadata={"help": "Number of heads in each layer to be applied guided attention loss, if set -1, all of the heads will be applied."},
+    )
+    modules_applied_guided_attn: Any = field(
+        default=("encoder-decoder",),
+        metadata={"help": "Module name list to be applied guided attention loss"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+
+
+class TexttoSpeechLoss(FairseqCriterion):
+    def __init__(
+        self,
+        task,
+        sentence_avg,
+        use_masking=True,
+        use_weighted_masking=False,
+        loss_type="L1",
+        bce_pos_weight=5.0,
+        bce_loss_lambda=1.0,
+        use_guided_attn_loss=False,
+        guided_attn_loss_sigma=0.4,
+        guided_attn_loss_lambda=1.0,
+        num_layers_applied_guided_attn=2,
+        num_heads_applied_guided_attn=2,
+        modules_applied_guided_attn=["encoder-decoder"],
+    ):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.use_masking = use_masking
+        self.use_weighted_masking = use_weighted_masking
+        self.loss_type = loss_type
+        self.bce_pos_weight = bce_pos_weight
+        self.bce_loss_lambda = bce_loss_lambda
+        self.use_guided_attn_loss = use_guided_attn_loss
+        self.guided_attn_loss_sigma = guided_attn_loss_sigma
+        self.guided_attn_loss_lambda = guided_attn_loss_lambda
+        # define loss function
+        self.criterion = Tacotron2Loss(
+            use_masking=use_masking,
+            use_weighted_masking=use_weighted_masking,
+            bce_pos_weight=bce_pos_weight,
+        )
+        if self.use_guided_attn_loss:
+            self.num_layers_applied_guided_attn = num_layers_applied_guided_attn
+            self.num_heads_applied_guided_attn = num_heads_applied_guided_attn
+            self.modules_applied_guided_attn = modules_applied_guided_attn
+        if self.use_guided_attn_loss:
+            self.attn_criterion = GuidedMultiHeadAttentionLoss(
+                sigma=guided_attn_loss_sigma,
+                alpha=guided_attn_loss_lambda,
+            )
+
+    def forward(self, model, sample):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        net_output = model(**sample["net_input"])
+        loss, l1_loss, l2_loss, bce_loss, enc_dec_attn_loss = self.compute_loss(model, net_output, sample)
+        # sample_size = (
+        #     sample["target"].size(0) if self.sentence_avg else sample["nframes"]
+        # )
+        sample_size = 1
+        logging_output = {
+            "loss": loss.item(),
+            "l1_loss": l1_loss.item(),
+            "l2_loss": l2_loss.item(),
+            "bce_loss": bce_loss.item(),
+            "sample_size": 1,
+            "ntokens": sample["ntokens"],
+            "nsentences": sample["target"].size(0),
+        }
+
+        if enc_dec_attn_loss is not None:
+            logging_output['enc_dec_attn_loss'] = enc_dec_attn_loss.item()
+
+        if hasattr(model, 'text_encoder_prenet'):
+            logging_output["encoder_alpha"] = model.text_encoder_prenet.encoder_prenet[-1].alpha.item()
+            logging_output["decoder_alpha"] = model.speech_decoder_prenet.decoder_prenet[-1].alpha.item()
+        elif hasattr(model, "speech_encoder_prenet"):
+            logging_output["decoder_alpha"] = model.speech_decoder_prenet.decoder_prenet[-1].alpha.item()
+        else:
+            if 'task' not in sample:
+                logging_output["encoder_alpha"] = model.encoder_prenet.encoder_prenet[-1].alpha.item()
+            logging_output["decoder_alpha"] = model.decoder_prenet.decoder_prenet[-1].alpha.item()
+
+        return loss, sample_size, logging_output
+
+    def compute_loss(self, model, net_output, sample):
+        before_outs, after_outs, logits, attn = net_output
+        labels = sample["labels"]
+        ys = sample["dec_target"]
+        olens = sample["dec_target_lengths"]
+        ilens = sample["src_lengths"]
+
+        # modifiy mod part of groundtruth
+        if model.reduction_factor > 1:
+            olens_in = olens.new([torch.div(olen, model.reduction_factor, rounding_mode='floor') for olen in olens])
+            olens = olens.new([olen - olen % model.reduction_factor for olen in olens])
+            max_olen = max(olens)
+            ys = ys[:, :max_olen]
+            labels = labels[:, :max_olen]
+            labels = torch.scatter(labels, 1, (olens - 1).unsqueeze(1), 1.0) # make sure at least one frame has 1
+            # labels[:, -1] = 1.0  
+        else:
+            olens_in = olens
+
+        # caluculate loss values
+        l1_loss, l2_loss, bce_loss = self.criterion(
+            after_outs, before_outs, logits, ys, labels, olens
+        )
+
+        # l1_loss = l1_loss / ys.size(2)
+        # l2_loss = l2_loss / ys.size(2)
+
+        if self.loss_type == "L1":
+            loss = l1_loss + self.bce_loss_lambda * bce_loss if self.bce_loss_lambda > 0.0 else l1_loss
+        elif self.loss_type == "L2":
+            loss = l2_loss + self.bce_loss_lambda * bce_loss if self.bce_loss_lambda > 0.0 else l2_loss
+        elif self.loss_type == "L1+L2":
+            loss = l1_loss + l2_loss + self.bce_loss_lambda * bce_loss if self.bce_loss_lambda > 0.0 else l1_loss + l2_loss
+        else:
+            raise ValueError("unknown --loss-type " + self.loss_type)
+
+        # calculate guided attention loss
+        enc_dec_attn_loss = None
+        if self.use_guided_attn_loss:
+            # calculate the input lengths of encoder, which is determined by encoder prenet
+            if hasattr(model, 'encoder_reduction_factor') and model.encoder_reduction_factor > 1:
+                ilens_in = ilens.new([ilen // model.encoder_reduction_factor for ilen in ilens])
+            else:
+                ilens_in = ilens
+            # work for speech to speech model's input
+            if "task_name" in sample and sample["task_name"] == "s2s":
+                m = None
+                if hasattr(model, 'encoder_prenet'):
+                    m = model.encoder_prenet
+                elif hasattr(model, 'speech_encoder_prenet'):
+                    m = model.speech_encoder_prenet
+                if m is not None and isinstance(m, SpeechEncoderPrenet):
+                    ilens_in = m.get_src_lengths(ilens_in)
+            # calculate for encoder-decoder
+            if "encoder-decoder" in self.modules_applied_guided_attn:
+                attn = [att_l[:, : self.num_heads_applied_guided_attn] for att_l in attn]
+                att_ws = torch.cat(attn, dim=1)  # (B, H*L, T_out, T_in)
+                enc_dec_attn_loss = self.attn_criterion(att_ws, ilens_in, olens_in)
+                loss = loss + enc_dec_attn_loss
+
+        return loss, l1_loss, l2_loss, bce_loss, enc_dec_attn_loss
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        l1_loss_sum = sum(log.get("l1_loss", 0) for log in logging_outputs)
+        l2_loss_sum = sum(log.get("l2_loss", 0) for log in logging_outputs)
+        bce_loss_sum = sum(log.get("bce_loss", 0) for log in logging_outputs)
+        sample_size = max(1, sum(log.get("sample_size", 0) for log in logging_outputs))
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size, sample_size, 1, round=5
+        )
+        encoder_alpha_sum = sum(log.get("encoder_alpha", 0) for log in logging_outputs)
+        decoder_alpha_sum = sum(log.get("decoder_alpha", 0) for log in logging_outputs)
+        ngpu = sum(log.get("ngpu", 0) for log in logging_outputs)
+
+        metrics.log_scalar(
+            "l1_loss", l1_loss_sum / sample_size, sample_size, 2, round=5
+        )
+        metrics.log_scalar(
+            "l2_loss", l2_loss_sum / sample_size, sample_size, 2, round=5
+        )
+        metrics.log_scalar(
+            "bce_loss", bce_loss_sum / sample_size, sample_size, 2, round=5
+        )
+        metrics.log_scalar(
+            "encoder_alpha", encoder_alpha_sum / sample_size, sample_size, round=5
+        )
+        metrics.log_scalar(
+            "decoder_alpha", decoder_alpha_sum / sample_size, sample_size, round=5
+        )
+
+        if "enc_dec_attn_loss" in logging_outputs[0]:
+            enc_dec_attn_loss_sum = sum(log.get("enc_dec_attn_loss", 0) for log in logging_outputs)
+            metrics.log_scalar(
+                "enc_dec_attn_loss", enc_dec_attn_loss_sum / sample_size, sample_size, round=8
+            )
+
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True
+
+class Tacotron2Loss(torch.nn.Module):
+    """Loss function module for Tacotron2."""
+
+    def __init__(
+        self, use_masking=True, use_weighted_masking=False, bce_pos_weight=20.0
+    ):
+        """Initialize Tactoron2 loss module.
+
+        Args:
+            use_masking (bool): Whether to apply masking
+                for padded part in loss calculation.
+            use_weighted_masking (bool):
+                Whether to apply weighted masking in loss calculation.
+            bce_pos_weight (float): Weight of positive sample of stop token.
+
+        """
+        super(Tacotron2Loss, self).__init__()
+        assert (use_masking != use_weighted_masking) or not use_masking
+        self.use_masking = use_masking
+        self.use_weighted_masking = use_weighted_masking
+
+        # define criterions
+        # reduction = "none" if self.use_weighted_masking else "sum"
+        reduction = "none" if self.use_weighted_masking else "mean"
+        self.l1_criterion = torch.nn.L1Loss(reduction=reduction)
+        self.mse_criterion = torch.nn.MSELoss(reduction=reduction)
+        self.bce_criterion = torch.nn.BCEWithLogitsLoss(
+            reduction=reduction, pos_weight=torch.tensor(bce_pos_weight)
+        )
+
+        # NOTE(kan-bayashi): register pre hook function for the compatibility
+        self._register_load_state_dict_pre_hook(self._load_state_dict_pre_hook)
+
+    def forward(self, after_outs, before_outs, logits, ys, labels, olens):
+        """Calculate forward propagation.
+
+        Args:
+            after_outs (Tensor): Batch of outputs after postnets (B, Lmax, odim).
+            before_outs (Tensor): Batch of outputs before postnets (B, Lmax, odim).
+            logits (Tensor): Batch of stop logits (B, Lmax).
+            ys (Tensor): Batch of padded target features (B, Lmax, odim).
+            labels (LongTensor): Batch of the sequences of stop token labels (B, Lmax).
+            olens (LongTensor): Batch of the lengths of each target (B,).
+
+        Returns:
+            Tensor: L1 loss value.
+            Tensor: Mean square error loss value.
+            Tensor: Binary cross entropy loss value.
+
+        """
+        # make mask and apply it
+        if self.use_masking:
+            masks = make_non_pad_mask(olens).unsqueeze(-1).to(ys.device)
+            ys = ys.masked_select(masks)
+            after_outs = after_outs.masked_select(masks)
+            before_outs = before_outs.masked_select(masks)
+            labels = labels.masked_select(masks[:, :, 0])
+            logits = logits.masked_select(masks[:, :, 0])
+
+        # calculate loss
+        l1_loss = self.l1_criterion(after_outs, ys) + self.l1_criterion(before_outs, ys)
+        mse_loss = self.mse_criterion(after_outs, ys) + self.mse_criterion(
+            before_outs, ys
+        )
+        bce_loss = self.bce_criterion(logits, labels)
+
+        # make weighted mask and apply it
+        if self.use_weighted_masking:
+            masks = make_non_pad_mask(olens).unsqueeze(-1).to(ys.device)
+            weights = masks.float() / masks.sum(dim=1, keepdim=True).float()
+            out_weights = weights.div(ys.size(0) * ys.size(2))
+            logit_weights = weights.div(ys.size(0))
+
+            # apply weight
+            l1_loss = l1_loss.mul(out_weights).masked_select(masks).sum()
+            mse_loss = mse_loss.mul(out_weights).masked_select(masks).sum()
+            bce_loss = (
+                bce_loss.mul(logit_weights.squeeze(-1))
+                .masked_select(masks.squeeze(-1))
+                .sum()
+            )
+
+        return l1_loss, mse_loss, bce_loss
+
+    def _load_state_dict_pre_hook(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        """Apply pre hook fucntion before loading state dict.
+
+        From v.0.6.1 `bce_criterion.pos_weight` param is registered as a parameter but
+        old models do not include it and as a result, it causes missing key error when
+        loading old model parameter. This function solve the issue by adding param in
+        state dict before loading as a pre hook function
+        of the `load_state_dict` method.
+
+        """
+        key = prefix + "bce_criterion.pos_weight"
+        if key not in state_dict:
+            state_dict[key] = self.bce_criterion.pos_weight
+
+class GuidedMultiHeadAttentionLoss(GuidedAttentionLoss):
+    """Guided attention loss function module for multi head attention.
+    Args:
+        sigma (float, optional): Standard deviation to control
+        how close attention to a diagonal.
+        alpha (float, optional): Scaling coefficient (lambda).
+        reset_always (bool, optional): Whether to always reset masks.
+    """
+
+    def forward(self, att_ws, ilens, olens):
+        """Calculate forward propagation.
+        Args:
+            att_ws (Tensor):
+                Batch of multi head attention weights (B, H, T_max_out, T_max_in).
+            ilens (LongTensor): Batch of input lenghts (B,).
+            olens (LongTensor): Batch of output lenghts (B,).
+        Returns:
+            Tensor: Guided attention loss value.
+        """
+        if self.guided_attn_masks is None:
+            self.guided_attn_masks = (
+                self._make_guided_attention_masks(ilens, olens)
+                .to(att_ws.device)
+                .unsqueeze(1)
+            )
+        if self.masks is None:
+            self.masks = self._make_masks(ilens, olens).to(att_ws.device).unsqueeze(1)
+        losses = self.guided_attn_masks * att_ws
+        loss = torch.mean(losses.masked_select(self.masks))
+        if self.reset_always:
+            self._reset_masks()
+
+        return self.alpha * loss
+
+    def _make_guided_attention_masks(self, ilens, olens):
+        n_batches = len(ilens)
+        max_ilen = max(ilens)
+        max_olen = max(olens)
+        guided_attn_masks = torch.zeros((n_batches, max_olen, max_ilen), device=olens.device)
+        for idx, (ilen, olen) in enumerate(zip(ilens, olens)):
+            guided_attn_masks[idx, :olen, :ilen] = self._make_guided_attention_mask(
+                ilen, olen, self.sigma
+            )
+        return guided_attn_masks
+
+    @staticmethod
+    def _make_guided_attention_mask(ilen, olen, sigma):
+        grid_x, grid_y = torch.meshgrid(torch.arange(olen, device=olen.device), torch.arange(ilen, device=olen.device))
+        grid_x, grid_y = grid_x.float(), grid_y.float()
+        return 1.0 - torch.exp(
+            -((grid_y / ilen - grid_x / olen) ** 2) / (2 * (sigma**2))
+        )
+
+    @staticmethod
+    def _make_masks(ilens, olens):
+        in_masks = make_non_pad_mask(ilens).to(ilens.device)  # (B, T_in)
+        out_masks = make_non_pad_mask(olens).to(olens.device)  # (B, T_out)
+        return out_masks.unsqueeze(-1) & in_masks.unsqueeze(-2)  # (B, T_out, T_in)

artst/data/multitask_dataset.py

@@ -0,0 +1,263 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import bisect
+
+import logging
+import numpy as np
+from torch.utils.data.dataloader import default_collate
+from fairseq.data import data_utils
+
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+logger = logging.getLogger(__name__)
+
+class MultitaskDataset(FairseqDataset):
+    @staticmethod
+    def cumsum(sequence):
+        r, s = [], 0
+        for e in sequence:
+            curr_len = len(e)
+            r.append(curr_len + s)
+            s += curr_len
+        return r
+
+    def __init__(self, datasets, sample_ratios=1, batch_ratio=None):
+        super(MultitaskDataset, self).__init__()
+        assert len(datasets) > 0, "datasets should not be an empty iterable"
+        self.datasets = list(datasets)
+        if isinstance(sample_ratios, int):
+            sample_ratios = [sample_ratios] * len(self.datasets)
+            if batch_ratio is not None:
+                logger.info('batch ratio is ' + str(batch_ratio))
+                self.batch_ratio = batch_ratio
+            else:
+                self.batch_ratio = None
+        else:
+            logger.info('set sample ratio to ' + str(sample_ratios))
+            if batch_ratio is not None:
+                logger.info('batch ratio is ' + str(batch_ratio))
+                self.batch_ratio = batch_ratio
+            else:
+                self.batch_ratio = None
+        self.sample_ratios = sample_ratios
+        self._ordered_indices = None
+        self._update_size()
+
+    def __len__(self):
+        return self.cumulative_sizes[-1]
+
+    def __getitem__(self, idx):
+        dataset_idx, sample_idx = self._get_dataset_and_sample_index(idx)
+        sample = self.datasets[dataset_idx][sample_idx]
+        if isinstance(sample, dict):
+            sample["dataset_idx"] = dataset_idx
+        else:
+            sample = sample + (dataset_idx,)
+        return sample
+
+    def _update_size(self):
+        self.cumulative_sizes = self.cumsum(self.datasets)
+        self.real_sizes = [len(d) for d in self.datasets]
+
+    def _get_dataset_and_sample_index(self, idx: int):
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        sample_idx = sample_idx % self.real_sizes[dataset_idx]
+        return dataset_idx, sample_idx
+
+    def collater(self, samples, **extra_args):
+        # For now only supports datasets with same underlying collater implementations
+        if samples is not None and len(samples) > 0:
+            if isinstance(samples[0], dict):
+                dataset_idx = samples[0]["dataset_idx"]
+            else:
+                dataset_idx = samples[0][-1]
+                samples = [sample[:-1] for sample in samples]
+        else:
+            dataset_idx = 0
+
+        if hasattr(self.datasets[dataset_idx], "collater"):
+            return self.datasets[dataset_idx].collater(samples, **extra_args)
+        else:
+            return default_collate(samples, **extra_args)
+
+    def size(self, idx: int):
+        """
+        Return an example's size as a float or tuple.
+        """
+        dataset_idx, sample_idx = self._get_dataset_and_sample_index(idx)
+        return self.datasets[dataset_idx].size(sample_idx)
+
+    def num_tokens(self, index: int):
+        return np.max(self.size(index))
+
+    def attr(self, attr: str, index: int):
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, index)
+        return getattr(self.datasets[dataset_idx], attr, None)
+
+    @property
+    def sizes(self):
+        _dataset_sizes = []
+        for ds in self.datasets:
+            if isinstance(ds.sizes, np.ndarray):
+                _dataset_sizes.append(ds.sizes)
+            else:
+                # Only support underlying dataset with single size array.
+                assert isinstance(ds.sizes, list)
+                _dataset_sizes.append(ds.sizes[0])
+        return np.concatenate(_dataset_sizes)
+
+    @property
+    def supports_prefetch(self):
+        return all(d.supports_prefetch for d in self.datasets)
+
+    def ordered_indices(self):
+        # ordered_indices = []
+        # for i, dataset in enumerate(self.datasets):
+        #     indice = dataset.ordered_indices()
+        #     ordered_indices.append(indice)
+        if self._ordered_indices is None:
+            # Call the underlying dataset's ordered_indices() here, so that we
+            # get the same random ordering as we would have from using the
+            # underlying sub-datasets directly.
+            self._ordered_indices = [
+                dataset.ordered_indices()
+                for dataset in self.datasets
+            ]
+        return np.arange(len(self))
+
+    def prefetch(self, indices):
+        frm = 0
+        for to, ds in zip(self.cumulative_sizes, self.datasets):
+            real_size = len(ds)
+            if getattr(ds, "supports_prefetch", False):
+                ds.prefetch([(i - frm) % real_size for i in indices if frm <= i < to])
+            frm = to
+
+    def batch_by_size(
+        self,
+        indices,
+        max_tokens=None,
+        max_sentences=None,
+        required_batch_size_multiple=1,
+    ):
+        if not hasattr(self, "max_tokens"):
+            self.max_tokens = max_tokens
+        if not hasattr(self, "max_sentences"):
+            self.max_sentences = max_sentences
+        if not hasattr(self, "required_batch_size_multiple"):
+            self.required_batch_size_multiple = required_batch_size_multiple
+        batch_samplers = []
+        for i, dataset in enumerate(self.datasets):
+            batch_sampler = dataset.batch_by_size(
+                self._ordered_indices[i],
+                max_tokens=max_tokens if self.batch_ratio is None else max_tokens * self.batch_ratio[i],
+                max_sentences=max_sentences,
+                required_batch_size_multiple=required_batch_size_multiple,
+            )
+            if i > 0:
+                for batch in batch_sampler:
+                    batch += self.cumulative_sizes[i - 1]
+            if self.sample_ratios[i] != 1.0:
+                batch_sampler = np.array(batch_sampler)
+                batch_sampler = np.random.choice(batch_sampler, int(len(batch_sampler) * self.sample_ratios[i]))
+                batch_sampler = list(batch_sampler)
+            logger.info('Adjust batch by ratio ' + str(self.sample_ratios[i]) + ' and the number of batch is ' + str(int(len(batch_sampler))) + ' for dataset ' + str(i))
+            batch_samplers.extend(batch_sampler)
+        return batch_samplers
+
+    def filter_indices_by_size(self, indices, max_positions):
+        """
+        Filter each sub-dataset independently, then update the round robin to work
+        on the filtered sub-datasets.
+        """
+        if not hasattr(self, "max_positions"):
+            self.max_positions = max_positions
+        ignored_some = False
+        for i in range(len(self.datasets)):
+            # ignored = []
+            self._ordered_indices[i], ignored = self.datasets[i].filter_indices_by_size(
+                self._ordered_indices[i], self.max_positions[i]
+            )
+            if len(ignored) > 0:
+                ignored_some = True
+                logger.warning(
+                    f"{len(ignored)} samples from {i} have invalid sizes and will be skipped, "
+                    f"max_positions={self.max_positions[i]}, first few sample ids={ignored[:10]}"
+                )
+
+        logger.info('update dataset size')
+        self._update_size()
+
+        # Since we are modifying in place the _ordered_indices,
+        # it's not possible anymore to return valid ignored indices.
+        # Hopefully the extra debug information print above should be enough to debug.
+        # Ideally we would receive ignore_invalid_inputs so that we could have
+        # a proper error message.
+        return (np.arange(len(self)), [0] if ignored_some else [])
+
+    @property
+    def can_reuse_epoch_itr_across_epochs(self):
+        return all(d.can_reuse_epoch_itr_across_epochs for d in self.datasets)
+
+    def set_epoch(self, epoch):
+        super().set_epoch(epoch)
+        for ds in self.datasets:
+            if hasattr(ds, "set_epoch"):
+                ds.set_epoch(epoch)
+
+    def shuffle_batches(self, batches, seed):
+        logger.info("shuffle batches")
+        new_batches_fromlist = []
+        new_batches_notlist = []
+        new_batches = []
+        with data_utils.numpy_seed(seed):
+            np.random.shuffle(batches)
+            for batch in batches:
+                if isinstance(batch, list):
+                    # np.random.shuffle(batch)
+                    new_batches_fromlist.append(batch)
+                else:
+                    new_batches_notlist.append(batch)
+            logger.info("Get " + str(len(new_batches_fromlist)) + " chunk from speech sides")
+            logger.info("Get " + str(sum([len(batch_list) for batch_list in new_batches_fromlist])) + " batches from speech sides")
+            logger.info("Get " + str(len(new_batches_notlist)) + " batches from text sides")
+            if len(new_batches_fromlist) == 0:
+                return new_batches_notlist
+            st_ratio = int(len(new_batches_notlist) / len(new_batches_fromlist))
+            logger.info("Get st_ratio " + str(st_ratio))
+            last_idx = 0
+            for i in range(len(new_batches_fromlist)):
+                if i == len(new_batches_fromlist) - 1:
+                    new_batches_fromlist[i].extend(new_batches_notlist[last_idx:])
+                else:
+                    new_batches_fromlist[i].extend(new_batches_notlist[last_idx : last_idx + st_ratio])
+                np.random.shuffle(new_batches_fromlist[i])
+                new_batches.extend(new_batches_fromlist[i])
+                last_idx = last_idx + st_ratio
+        logger.info("Finish shuffle")
+        return new_batches
+
+    def reset_batch_sampler(self):
+        logger.info("reset batch sampler")
+        self._ordered_indices = [
+            self.datasets[i].ordered_indices()
+            for i in range(len(self.datasets))
+        ]
+        self.filter_indices_by_size(None, None)
+
+        batch_samplers = self.batch_by_size(
+            None,
+            self.max_tokens,
+            self.max_sentences,
+            self.required_batch_size_multiple
+        )
+        return batch_samplers

artst/data/speech_dataset.py

@@ -0,0 +1,475 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import itertools
+import logging
+import os
+import sys
+from typing import Any, List, Optional, Union
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+import librosa
+from fairseq.data.audio.speech_to_text_dataset import get_features_or_waveform
+from fairseq.data import data_utils
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+logger = logging.getLogger(__name__)
+
+def _collate_frames(
+    frames: List[torch.Tensor], is_audio_input: bool = False
+):
+    """
+    Convert a list of 2D frames into a padded 3D tensor
+    Args:
+        frames (list): list of 2D frames of size L[i]*f_dim. Where L[i] is
+            length of i-th frame and f_dim is static dimension of features
+    Returns:
+        3D tensor of size len(frames)*len_max*f_dim where len_max is max of L[i]
+    """
+    max_len = max(frame.size(0) for frame in frames)
+    if is_audio_input:
+        out = frames[0].new_zeros((len(frames), max_len))
+    else:
+        out = frames[0].new_zeros((len(frames), max_len, frames[0].size(1)))
+    for i, v in enumerate(frames):
+        out[i, : v.size(0)] = v
+    return out
+
+def add_first_frame_and_remove_last_frame(ys):
+    ys_in = torch.cat(
+        [ys.new_zeros((ys.shape[0], 1, ys.shape[2])), ys[:, :-1]], dim=1
+    )
+    return ys_in
+
+def load_audio(manifest_path, max_keep, min_keep):
+    n_long, n_short = 0, 0
+    names, inds, sizes, spk_embeds = [], [], [], []
+    with open(manifest_path) as f:
+        root = f.readline().strip()
+        for ind, line in enumerate(f):
+            items = line.strip().split("\t")
+            assert len(items) == 3, line
+            sz = int(items[1])
+            if min_keep is not None and sz < min_keep:
+                n_short += 1
+            elif max_keep is not None and sz > max_keep:
+                n_long += 1
+            else:
+                names.append(items[0])
+                spk_embeds.append(items[2])
+                inds.append(ind)
+                sizes.append(sz)
+    tot = ind + 1
+    logger.info(
+        (
+            f"max_keep={max_keep}, min_keep={min_keep}, "
+            f"loaded {len(names)}, skipped {n_short} short and {n_long} long, "
+            f"longest-loaded={max(sizes)}, shortest-loaded={min(sizes)}"
+        )
+    )
+    return root, names, inds, tot, sizes, spk_embeds
+
+
+def load_label(label_path, inds, tot):
+    with open(label_path) as f:
+        labels = [line.rstrip() for line in f]
+        assert (
+            len(labels) == tot
+        ), f"number of labels does not match ({len(labels)} != {tot})"
+        labels = [labels[i] for i in inds]
+    return labels
+
+
+def load_label_offset(label_path, inds, tot):
+    with open(label_path) as f:
+        code_lengths = [len(line.encode("utf-8")) for line in f]
+        assert (
+            len(code_lengths) == tot
+        ), f"number of labels does not match ({len(code_lengths)} != {tot})"
+        offsets = list(itertools.accumulate([0] + code_lengths))
+        offsets = [(offsets[i], offsets[i + 1]) for i in inds]
+    return offsets
+
+
+def verify_label_lengths(
+    audio_sizes,
+    audio_rate,
+    label_path,
+    label_rate,
+    inds,
+    tot,
+    tol=0.1,  # tolerance in seconds
+):
+    if label_rate < 0:
+        logger.info(f"{label_path} is sequence label. skipped")
+        return
+
+    with open(label_path) as f:
+        lengths = [len(line.rstrip().split()) for line in f]
+        assert len(lengths) == tot
+        lengths = [lengths[i] for i in inds]
+    num_invalid = 0
+    for i, ind in enumerate(inds):
+        dur_from_audio = audio_sizes[i] / audio_rate
+        dur_from_label = lengths[i] / label_rate
+        if abs(dur_from_audio - dur_from_label) > tol:
+            logger.warning(
+                (
+                    f"audio and label duration differ too much "
+                    f"(|{dur_from_audio} - {dur_from_label}| > {tol}) "
+                    f"in line {ind+1} of {label_path}. Check if `label_rate` "
+                    f"is correctly set (currently {label_rate}). "
+                    f"num. of samples = {audio_sizes[i]}; "
+                    f"label length = {lengths[i]}"
+                )
+            )
+            num_invalid += 1
+    if num_invalid > 0:
+        logger.warning(
+            f"total {num_invalid} (audio, label) pairs with mismatched lengths"
+        )
+
+
+def logmelfilterbank(
+    audio,
+    sampling_rate,
+    fft_size=1024,
+    hop_size=256,
+    win_length=None,
+    window="hann",
+    num_mels=80,
+    fmin=80,
+    fmax=7600,
+    eps=1e-10,
+):
+    """Compute log-Mel filterbank feature. 
+    (https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/bin/preprocess.py)
+
+    Args:
+        audio (ndarray): Audio signal (T,).
+        sampling_rate (int): Sampling rate.
+        fft_size (int): FFT size.
+        hop_size (int): Hop size.
+        win_length (int): Window length. If set to None, it will be the same as fft_size.
+        window (str): Window function type.
+        num_mels (int): Number of mel basis.
+        fmin (int): Minimum frequency in mel basis calculation.
+        fmax (int): Maximum frequency in mel basis calculation.
+        eps (float): Epsilon value to avoid inf in log calculation.
+
+    Returns:
+        ndarray: Log Mel filterbank feature (#frames, num_mels).
+
+    """
+    # get amplitude spectrogram
+    x_stft = librosa.stft(audio, n_fft=fft_size, hop_length=hop_size,
+                          win_length=win_length, window=window, pad_mode="reflect")
+    spc = np.abs(x_stft).T  # (#frames, #bins)
+
+    # get mel basis
+    fmin = 0 if fmin is None else fmin
+    fmax = sampling_rate / 2 if fmax is None else fmax
+    mel_basis = librosa.filters.mel(sr=sampling_rate, n_fft=fft_size, n_mels=num_mels, fmin=fmin, fmax=fmax)
+
+    return np.log10(np.maximum(eps, np.dot(spc, mel_basis.T)))
+
+
+class SpeechPretrainDataset(FairseqDataset):
+    def __init__(
+        self,
+        manifest_path: str,
+        sample_rate: float,
+        label_paths: List[str],
+        label_rates: Union[List[float], float],  # -1 for sequence labels
+        pad_list: List[str],
+        eos_list: List[str],
+        label_processors: Optional[List[Any]] = None,
+        max_keep_sample_size: Optional[int] = None,
+        min_keep_sample_size: Optional[int] = None,
+        max_sample_size: Optional[int] = None,
+        shuffle: bool = True,
+        pad_audio: bool = False,
+        normalize: bool = False,
+        store_labels: bool = True,
+        random_crop: bool = False,
+        single_target: bool = False,
+        reduction_factor: int = 1,
+    ):
+        self.audio_root, self.audio_names, inds, tot, self.sizes, self.spk_embeds = load_audio(
+            manifest_path, max_keep_sample_size, min_keep_sample_size
+        )
+        self.sample_rate = sample_rate
+        self.shuffle = shuffle
+        self.random_crop = random_crop
+
+        self.num_labels = len(label_paths)
+        self.pad_list = pad_list
+        self.eos_list = eos_list
+        self.label_processors = label_processors
+        self.single_target = single_target
+        self.label_rates = (
+            [label_rates for _ in range(len(label_paths))]
+            if isinstance(label_rates, float)
+            else label_rates
+        )
+        self.store_labels = store_labels
+        if store_labels:
+            self.label_list = [load_label(p, inds, tot) for p in label_paths]
+        else:
+            self.label_paths = label_paths
+            self.label_offsets_list = [
+                load_label_offset(p, inds, tot) for p in label_paths
+            ]
+        assert label_processors is None or len(label_processors) == self.num_labels
+        for label_path, label_rate in zip(label_paths, self.label_rates):
+            verify_label_lengths(
+                self.sizes, sample_rate, label_path, label_rate, inds, tot
+            )
+
+        self.max_sample_size = (
+            max_sample_size if max_sample_size is not None else sys.maxsize
+        )
+        self.pad_audio = pad_audio
+        self.normalize = normalize
+        self.reduction_factor = reduction_factor
+        logger.info(
+            f"pad_audio={pad_audio}, random_crop={random_crop}, reduction_factor={reduction_factor}, "
+            f"normalize={normalize}, max_sample_size={self.max_sample_size}"
+        )
+
+    def get_audio(self, index):
+        import soundfile as sf
+
+        wav_path = os.path.join(self.audio_root, self.audio_names[index])
+        wav, cur_sample_rate = sf.read(wav_path)
+        wav = torch.from_numpy(wav).float()
+        fbank = logmelfilterbank(
+            wav.view(-1).cpu().numpy(), 16000
+        )
+        fbank = torch.from_numpy(fbank).float()
+        wav = self.postprocess(wav, cur_sample_rate)
+        return wav, fbank
+
+    def get_label(self, index, label_idx):
+        if self.store_labels:
+            label = self.label_list[label_idx][index]
+        else:
+            with open(self.label_paths[label_idx]) as f:
+                offset_s, offset_e = self.label_offsets_list[label_idx][index]
+                f.seek(offset_s)
+                label = f.read(offset_e - offset_s)
+
+        if self.label_processors is not None:
+            label = self.label_processors[label_idx](label)
+        return label
+
+    def get_labels(self, index):
+        return [self.get_label(index, i) for i in range(self.num_labels)]
+
+    def __getitem__(self, index):
+        wav, fbank = self.get_audio(index)
+        labels = self.get_labels(index)
+        spkembs = get_features_or_waveform(
+            os.path.join(self.audio_root, self.spk_embeds[index])
+        )
+        spkembs = torch.from_numpy(spkembs).float()
+        return {"id": index, "source": wav, "target": fbank, "label_list": labels, 'spkembs': spkembs}
+
+    def __len__(self):
+        return len(self.sizes)
+
+    def crop_to_max_size(self, wav, target_size):
+        size = len(wav)
+        diff = size - target_size
+        if diff <= 0:
+            return wav, 0
+
+        start, end = 0, target_size
+        if self.random_crop:
+            start = np.random.randint(0, diff + 1)
+            end = size - diff + start
+        return wav[start:end], start
+
+    def collater(self, samples):
+        # target = max(sizes) -> random_crop not used
+        # target = max_sample_size -> random_crop used for long
+        samples = [s for s in samples if s["source"] is not None]
+        if len(samples) == 0:
+            return {}
+
+        audios = [s["source"] for s in samples]
+        audio_sizes = [len(s) for s in audios]
+
+        fbanks = [s["target"] for s in samples]
+        fbank_sizes = [len(s) for s in fbanks]
+        
+        if self.pad_audio:
+            audio_size = min(max(audio_sizes), self.max_sample_size)
+        else:
+            audio_size = min(min(audio_sizes), self.max_sample_size)
+        collated_audios, padding_mask, audio_starts = self.collater_audio(
+            audios, audio_size
+        )
+
+        collated_fbanks = []
+        collated_audios_size = []
+        for i in range(len(fbanks)):
+            fbank_start = int(audio_starts[i] / (audio_sizes[i] / fbank_sizes[i]))
+            fbank_size = int(audio_size / (audio_sizes[i] / fbank_sizes[i]))
+            fbank_end = min(fbank_start + fbank_size, fbank_sizes[i])
+            collated_fbanks.append(fbanks[i][fbank_start : fbank_end])
+            collated_audios_size.append(audio_size)
+        collated_fbanks_size = [len(s) for s in collated_fbanks]
+        collated_fbanks = _collate_frames(collated_fbanks)
+        collated_fbanks_size = torch.tensor(collated_fbanks_size, dtype=torch.long)
+
+        # thin out frames for reduction factor (B, Lmax, odim) ->  (B, Lmax//r, odim)
+        if self.reduction_factor > 1:
+            collated_fbanks_in = collated_fbanks[:, self.reduction_factor - 1 :: self.reduction_factor]
+            collated_fbanks_size_in = collated_fbanks_size.new([torch.div(olen, self.reduction_factor, rounding_mode='floor') for olen in collated_fbanks_size])
+        else:
+            collated_fbanks_in, collated_fbanks_size_in = collated_fbanks, collated_fbanks_size
+
+        prev_output_tokens = torch.cat(
+            [collated_fbanks_in.new_zeros((collated_fbanks_in.shape[0], 1, collated_fbanks_in.shape[2])), collated_fbanks_in[:, :-1]], dim=1
+        )
+
+	    # make labels for stop prediction
+        labels = collated_fbanks.new_zeros(collated_fbanks.size(0), collated_fbanks.size(1))
+        for i, l in enumerate(fbank_sizes):
+            labels[i, l - 1 :] = 1.0
+
+        spkembs = _collate_frames([s["spkembs"] for s in samples], is_audio_input=True)
+
+        targets_by_label = [
+            [s["label_list"][i] for s in samples] for i in range(self.num_labels)
+        ]
+        targets_list, lengths_list, ntokens_list = self.collater_label(
+            targets_by_label, audio_size, audio_starts
+        )
+
+        net_input = {
+            "source": collated_audios, 
+            "padding_mask": padding_mask, 
+            "prev_output_tokens": prev_output_tokens,
+            "spkembs": spkembs,
+            "tgt_lengths": collated_fbanks_size_in,
+        }
+
+        batch = {
+            "id": torch.LongTensor([s["id"] for s in samples]),
+            "net_input": net_input,
+            "labels": labels,
+            "dec_target": collated_fbanks,
+            "dec_target_lengths": collated_fbanks_size,
+            "src_lengths": collated_audios_size,
+            "task_name": 'speech_pretrain',
+        }
+
+        if self.single_target:
+            batch["target_lengths"] = lengths_list[0]
+            batch["ntokens"] = ntokens_list[0]
+            batch["target"] = targets_list[0]
+        else:
+            batch["target_lengths_list"] = lengths_list
+            batch["ntokens_list"] = ntokens_list
+            batch["target_list"] = targets_list
+        return batch
+
+    def collater_audio(self, audios, audio_size):
+        collated_audios = audios[0].new_zeros(len(audios), audio_size)
+        padding_mask = (
+            torch.BoolTensor(collated_audios.shape).fill_(False)
+            # if self.pad_audio else None
+        )
+        audio_starts = [0 for _ in audios]
+        for i, audio in enumerate(audios):
+            diff = len(audio) - audio_size
+            if diff == 0:
+                collated_audios[i] = audio
+            elif diff < 0:
+                assert self.pad_audio
+                collated_audios[i] = torch.cat([audio, audio.new_full((-diff,), 0.0)])
+                padding_mask[i, diff:] = True
+            else:
+                collated_audios[i], audio_starts[i] = self.crop_to_max_size(
+                    audio, audio_size
+                )
+        return collated_audios, padding_mask, audio_starts
+
+    def collater_frm_label(self, targets, audio_size, audio_starts, label_rate, pad):
+        assert label_rate > 0
+        s2f = label_rate / self.sample_rate
+        frm_starts = [int(round(s * s2f)) for s in audio_starts]
+        frm_size = int(round(audio_size * s2f))
+        if not self.pad_audio:
+            rem_size = [len(t) - s for t, s in zip(targets, frm_starts)]
+            frm_size = min(frm_size, *rem_size)
+        targets = [t[s : s + frm_size] for t, s in zip(targets, frm_starts)]
+        logger.debug(f"audio_starts={audio_starts}")
+        logger.debug(f"frame_starts={frm_starts}")
+        logger.debug(f"frame_size={frm_size}")
+
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_seq_label(self, targets, pad):
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_label(self, targets_by_label, audio_size, audio_starts):
+        targets_list, lengths_list, ntokens_list = [], [], []
+        itr = zip(targets_by_label, self.label_rates, self.pad_list)
+        for targets, label_rate, pad in itr:
+            if label_rate == -1.0:
+                targets, lengths, ntokens = self.collater_seq_label(targets, pad)
+            else:
+                targets, lengths, ntokens = self.collater_frm_label(
+                    targets, audio_size, audio_starts, label_rate, pad
+                )
+            targets_list.append(targets)
+            lengths_list.append(lengths)
+            ntokens_list.append(ntokens)
+        return targets_list, lengths_list, ntokens_list
+
+    def num_tokens(self, index):
+        return self.size(index)
+
+    def size(self, index):
+        if self.pad_audio:
+            return self.sizes[index]
+        return min(self.sizes[index], self.max_sample_size)
+
+    def ordered_indices(self):
+        if self.shuffle:
+            order = [np.random.permutation(len(self))]
+        else:
+            order = [np.arange(len(self))]
+
+        order.append(self.sizes)
+        return np.lexsort(order)[::-1]
+
+    def postprocess(self, wav, cur_sample_rate):
+        if wav.dim() == 2:
+            wav = wav.mean(-1)
+        assert wav.dim() == 1, wav.dim()
+
+        if cur_sample_rate != self.sample_rate:
+            raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")
+
+        if self.normalize:
+            with torch.no_grad():
+                wav = F.layer_norm(wav, wav.shape)
+        return wav

artst/data/speech_to_class_dataset.py

@@ -0,0 +1,260 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import logging
+import os
+from typing import Any, List, Optional
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from fairseq.data import data_utils, Dictionary
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+logger = logging.getLogger(__name__)
+
+
+def load_audio(manifest_path, max_keep, min_keep):
+    """manifest tsv: wav_path, wav_nframe, wav_class
+
+    Args
+        manifest_path: str
+        max_keep: int
+        min_keep: int
+    
+    Return
+        root, names, inds, tot, sizes, classes
+    """
+    n_long, n_short = 0, 0
+    names, inds, sizes, classes = [], [], [], []
+    with open(manifest_path) as f:
+        root = f.readline().strip()
+        for ind, line in enumerate(f):
+            items = line.strip().split("\t")
+            assert len(items) >= 2, line
+            sz = int(items[1])
+            if min_keep is not None and sz < min_keep:
+                n_short += 1
+            elif max_keep is not None and sz > max_keep:
+                n_long += 1
+            else:
+                names.append(items[0])
+                if len(items) > 2:
+                    classes.append(items[2])
+                inds.append(ind)
+                sizes.append(sz)
+    tot = ind + 1
+    logger.info(
+        (
+            f"max_keep={max_keep}, min_keep={min_keep}, "
+            f"loaded {len(names)}, skipped {n_short} short and {n_long} long, "
+            f"longest-loaded={max(sizes)}, shortest-loaded={min(sizes)}"
+        )
+    )
+    if len(classes) == 0:
+        logger.warn("no classes loaded only if inference")
+    return root, names, inds, tot, sizes, classes
+
+
+def sample_from_feature(x: np.ndarray, max_segment_length: int = 300):
+    """Load a segment within 300-400/51200-76800 frames or the corresponding samples from a utterance.
+
+    Args:
+        x (np.ndarray): feature or waveform (frames[, features]), e.g., log mel filter bank or waveform
+        max_segment_length (int, optional): maximum segment length. Defaults to 400.
+
+    Returns:
+        np.ndarray: segmented features
+    """
+    if len(x) <= max_segment_length:
+        return x
+    start = np.random.randint(0, x.shape[0] - max_segment_length)
+    return x[start: start + max_segment_length]
+
+
+class SpeechToClassDataset(FairseqDataset):
+    def __init__(
+        self,
+        manifest_path: str,
+        sample_rate: float,
+        label_processors: Optional[List[Any]] = None,
+        max_keep_sample_size: Optional[int] = None,
+        min_keep_sample_size: Optional[int] = None,
+        shuffle: bool = True,
+        normalize: bool = False,
+        tgt_dict: Optional[Dictionary] = None,
+        max_length: Optional[int] = None
+    ):
+        self.audio_root, self.audio_names, inds, tot, self.wav_sizes, self.wav_classes = load_audio(
+            manifest_path, max_keep_sample_size, min_keep_sample_size
+        )
+        self.sample_rate = sample_rate
+        self.shuffle = shuffle
+
+        self.label_processors = label_processors
+
+        self.normalize = normalize
+        self.tgt_dict = tgt_dict
+        self.max_length = max_length
+        logger.info(
+            f"max_length={max_length}, normalize={normalize}"
+        )
+
+    def get_audio(self, index):
+        import soundfile as sf
+
+        wav_path = os.path.join(self.audio_root, self.audio_names[index])
+        wav, cur_sample_rate = sf.read(wav_path)
+        if self.max_length is not None:
+            wav = sample_from_feature(wav, self.max_length)
+        wav = torch.from_numpy(wav).float()
+        wav = self.postprocess(wav, cur_sample_rate)
+        return wav
+
+    def get_label(self, index):
+        label = self.wav_classes[index]
+
+        if self.label_processors is not None:
+            label = self.label_processors(label)
+        return label
+
+    def __getitem__(self, index):
+        wav = self.get_audio(index)
+        label = None
+        if len(self.wav_classes) == len(self.audio_names):
+            label = self.get_label(index)
+        return {"id": index, "source": wav, "label": label}
+
+    def __len__(self):
+        return len(self.wav_sizes)
+
+    def collater(self, samples):
+        samples = [s for s in samples if s["source"] is not None]
+        if len(samples) == 0:
+            return {}
+
+        audios = [s["source"] for s in samples]
+        audio_sizes = [len(s) for s in audios]
+
+        audio_size = max(audio_sizes)
+        collated_audios, padding_mask = self.collater_audio(
+            audios, audio_size
+        )
+
+        decoder_label = None
+        decoder_target = None
+        decoder_target_lengths = None
+        if samples[0]["label"] is not None:
+            targets_by_label = [
+                [s["label"] for s in samples]
+            ]
+            targets_list, lengths_list, ntokens_list = self.collater_label(targets_by_label)
+
+            decoder_label = [
+                (targets_list[0][i, :lengths_list[0][i]]).long()
+                for i in range(targets_list[0].size(0))
+            ]
+
+            decoder_target = data_utils.collate_tokens(
+                decoder_label,
+                self.tgt_dict.pad(),
+                self.tgt_dict.eos(),
+                left_pad=False,
+                move_eos_to_beginning=False,
+            )
+            decoder_target_lengths = torch.tensor(
+                [x.size(0) for x in decoder_label], dtype=torch.long
+            )
+        prev_output_tokens = data_utils.collate_tokens(
+            [torch.LongTensor([-1]) for _ in samples],
+            self.tgt_dict.pad(),
+            self.tgt_dict.eos(),
+            left_pad=False,
+            move_eos_to_beginning=True,
+        )
+
+        net_input = {
+            "source": collated_audios, 
+            "padding_mask": padding_mask,
+            "prev_output_tokens": prev_output_tokens,
+            "task_name": "s2c",
+        }
+        batch = {
+            "id": torch.LongTensor([s["id"] for s in samples]),
+            "net_input": net_input,
+            "target": decoder_target,
+            "target_lengths": decoder_target_lengths,
+            "task_name": "s2c",
+            "ntokens": len(samples),
+        }
+
+        return batch
+
+    def collater_audio(self, audios, audio_size):
+        collated_audios = audios[0].new_zeros(len(audios), audio_size)
+        padding_mask = (
+            torch.BoolTensor(collated_audios.shape).fill_(False)
+        )
+        for i, audio in enumerate(audios):
+            diff = len(audio) - audio_size
+            if diff == 0:
+                collated_audios[i] = audio
+            elif diff < 0:
+                collated_audios[i] = torch.cat([audio, audio.new_full((-diff,), 0.0)])
+                padding_mask[i, diff:] = True
+            else:
+                raise Exception("Diff should not be larger than 0")
+        return collated_audios, padding_mask
+
+    def collater_seq_label(self, targets, pad):
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_label(self, targets_by_label):
+        targets_list, lengths_list, ntokens_list = [], [], []
+        itr = zip(targets_by_label, [self.tgt_dict.pad()])
+        for targets, pad in itr:
+            targets, lengths, ntokens = self.collater_seq_label(targets, pad)
+            targets_list.append(targets)
+            lengths_list.append(lengths)
+            ntokens_list.append(ntokens)
+        return targets_list, lengths_list, ntokens_list
+
+    def num_tokens(self, index):
+        return self.size(index)
+
+    def size(self, index):
+        return self.wav_sizes[index]
+
+    @property
+    def sizes(self):
+        return np.array(self.wav_sizes)
+
+    def ordered_indices(self):
+        if self.shuffle:
+            order = [np.random.permutation(len(self))]
+        else:
+            order = [np.arange(len(self))]
+
+        order.append(self.wav_sizes)
+        return np.lexsort(order)[::-1]
+
+    def postprocess(self, wav, cur_sample_rate):
+        if wav.dim() == 2:
+            wav = wav.mean(-1)
+        assert wav.dim() == 1, wav.dim()
+
+        if cur_sample_rate != self.sample_rate:
+            raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")
+
+        if self.normalize:
+            with torch.no_grad():
+                wav = F.layer_norm(wav, wav.shape)
+        return wav

artst/data/speech_to_speech_dataset.py

@@ -0,0 +1,280 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import logging
+import os
+from typing import Any, List, Optional
+
+import librosa
+import numpy as np
+import torch
+import torch.nn.functional as F
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+logger = logging.getLogger(__name__)
+
+def _collate_frames(
+    frames: List[torch.Tensor], is_audio_input: bool = False
+):
+    """
+    Convert a list of 2D frames into a padded 3D tensor
+    Args:
+        frames (list): list of 2D frames of size L[i]*f_dim. Where L[i] is
+            length of i-th frame and f_dim is static dimension of features
+    Returns:
+        3D tensor of size len(frames)*len_max*f_dim where len_max is max of L[i]
+    """
+    max_len = max(frame.size(0) for frame in frames)
+    if is_audio_input:
+        out = frames[0].new_zeros((len(frames), max_len))
+    else:
+        out = frames[0].new_zeros((len(frames), max_len, frames[0].size(1)))
+    for i, v in enumerate(frames):
+        out[i, : v.size(0)] = v
+    return out
+
+def load_audio(manifest_path, max_keep, min_keep):
+    """manifest tsv: src_wav, src_nframe, tgt_wav, tgt_nframe, tgt_spkemb"""
+    n_long, n_short = 0, 0
+    src_names, tgt_names, inds, sizes, tgt_sizes, spk_embeds = [], [], [], [], [], []
+    with open(manifest_path) as f:
+        root = f.readline().strip()
+        for ind, line in enumerate(f):
+            items = line.strip().split("\t")
+            assert len(items) >= 2, line
+            sz = int(items[1])
+            if min_keep is not None and sz < min_keep:
+                n_short += 1
+            elif max_keep is not None and sz > max_keep:
+                n_long += 1
+            else:
+                src_names.append(items[0])
+                tgt_names.append(items[2])
+                tgt_sizes.append(items[3])
+                spk_embeds.append(items[4])
+                inds.append(ind)
+                sizes.append(sz)
+    tot = ind + 1
+    logger.info(
+        (
+            f"max_keep={max_keep}, min_keep={min_keep}, "
+            f"loaded {len(src_names)}, skipped {n_short} short and {n_long} long, "
+            f"longest-loaded={max(sizes)}, shortest-loaded={min(sizes)}"
+        )
+    )
+    return root, src_names, inds, tot, sizes, tgt_names, tgt_sizes, spk_embeds
+
+
+def logmelfilterbank(
+    audio,
+    sampling_rate,
+    fft_size=1024,
+    hop_size=256,
+    win_length=None,
+    window="hann",
+    num_mels=80,
+    fmin=80,
+    fmax=7600,
+    eps=1e-10,
+):
+    """Compute log-Mel filterbank feature. 
+    (https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/bin/preprocess.py)
+
+    Args:
+        audio (ndarray): Audio signal (T,).
+        sampling_rate (int): Sampling rate.
+        fft_size (int): FFT size.
+        hop_size (int): Hop size.
+        win_length (int): Window length. If set to None, it will be the same as fft_size.
+        window (str): Window function type.
+        num_mels (int): Number of mel basis.
+        fmin (int): Minimum frequency in mel basis calculation.
+        fmax (int): Maximum frequency in mel basis calculation.
+        eps (float): Epsilon value to avoid inf in log calculation.
+
+    Returns:
+        ndarray: Log Mel filterbank feature (#frames, num_mels).
+
+    """
+    # get amplitude spectrogram
+    x_stft = librosa.stft(audio, n_fft=fft_size, hop_length=hop_size,
+                          win_length=win_length, window=window, pad_mode="reflect")
+    spc = np.abs(x_stft).T  # (#frames, #bins)
+
+    # get mel basis
+    fmin = 0 if fmin is None else fmin
+    fmax = sampling_rate / 2 if fmax is None else fmax
+    mel_basis = librosa.filters.mel(sr=sampling_rate, n_fft=fft_size, n_mels=num_mels, fmin=fmin, fmax=fmax)
+
+    return np.log10(np.maximum(eps, np.dot(spc, mel_basis.T)))
+
+
+class SpeechToSpeechDataset(FairseqDataset):
+    def __init__(
+        self,
+        manifest_path: str,
+        sample_rate: float,
+        max_keep_sample_size: Optional[int] = None,
+        min_keep_sample_size: Optional[int] = None,
+        shuffle: bool = True,
+        normalize: bool = False,
+        reduction_factor: int = 1,
+    ):
+        self.audio_root, self.audio_names, inds, tot, self.wav_sizes, self.tgt_audios, self.tgt_sizes, self.tgt_spkembs = load_audio(
+            manifest_path, max_keep_sample_size, min_keep_sample_size
+        )
+        self.sample_rate = sample_rate
+        self.shuffle = shuffle
+
+        self.normalize = normalize
+        self.reduction_factor = reduction_factor
+        logger.info(
+            f"reduction_factor={reduction_factor}, normalize={normalize}"
+        )
+
+    def get_audio(self, index):
+        import soundfile as sf
+
+        wav_fbank = []
+        for name in [self.audio_names[index], self.tgt_audios[index]]:
+            wav_path = os.path.join(self.audio_root, name)
+            wav, cur_sample_rate = sf.read(wav_path)
+            wav = torch.from_numpy(wav).float()
+            fbank = logmelfilterbank(
+                wav.view(-1).cpu().numpy(), 16000
+            )
+            fbank = torch.from_numpy(fbank).float()
+            wav = self.postprocess(wav, cur_sample_rate)
+            wav_fbank.append(wav)
+            wav_fbank.append(fbank)
+        src_wav, src_fbank, tgt_wav, tgt_fbank = wav_fbank
+        return src_wav, src_fbank, tgt_wav, tgt_fbank
+
+    def __getitem__(self, index):
+        src_wav, src_fbank, tgt_wav, tgt_fbank = self.get_audio(index)
+        spkembs = np.load(os.path.join(self.audio_root, self.tgt_spkembs[index]))
+        spkembs = torch.from_numpy(spkembs).float()
+        name = self.audio_names[index].replace("/", ".").replace(".wav", "") + "-" + self.tgt_audios[index].replace("/", ".").replace(".wav", "") + ".wav"
+        return {"id": index, "source": src_wav, "target": tgt_fbank, "spkembs": spkembs, "audio_name": name, "tgt_name": self.tgt_audios[index]}
+    
+    def __len__(self):
+        return len(self.wav_sizes)
+
+    def collater(self, samples):
+        samples = [s for s in samples if s["source"] is not None]
+        if len(samples) == 0:
+            return {}
+
+        audios = [s["source"] for s in samples]
+        audio_sizes = [len(s) for s in audios]
+
+        audio_size = max(audio_sizes)
+        collated_audios, padding_mask = self.collater_audio(
+            audios, audio_size
+        )
+
+        fbanks = [s["target"] for s in samples]
+        fbank_sizes = [len(s) for s in fbanks]
+
+        collated_fbanks = _collate_frames(fbanks)
+        collated_fbanks_size = torch.tensor(fbank_sizes, dtype=torch.long)
+
+        # thin out frames for reduction factor (B, Lmax, odim) ->  (B, Lmax//r, odim)
+        if self.reduction_factor > 1:
+            collated_fbanks_in = collated_fbanks[:, self.reduction_factor - 1 :: self.reduction_factor]
+            collated_fbanks_size_in = collated_fbanks_size.new([torch.div(olen, self.reduction_factor, rounding_mode='floor') for olen in collated_fbanks_size])
+        else:
+            collated_fbanks_in, collated_fbanks_size_in = collated_fbanks, collated_fbanks_size
+
+        prev_output_tokens = torch.cat(
+            [collated_fbanks_in.new_zeros((collated_fbanks_in.shape[0], 1, collated_fbanks_in.shape[2])), collated_fbanks_in[:, :-1]], dim=1
+        )
+
+        # make labels for stop prediction
+        labels = collated_fbanks.new_zeros(collated_fbanks.size(0), collated_fbanks.size(1))
+        for i, l in enumerate(fbank_sizes):
+            labels[i, l - 1 :] = 1.0
+
+        spkembs = _collate_frames([s["spkembs"] for s in samples], is_audio_input=True)
+
+        net_input = {
+            "source": collated_audios,
+            "padding_mask": padding_mask,
+            "prev_output_tokens": prev_output_tokens,
+            "tgt_lengths": collated_fbanks_size_in,
+            "spkembs": spkembs,
+            "task_name": "s2s",
+        }
+        batch = {
+            "id": torch.LongTensor([s["id"] for s in samples]),
+            "name": [s["audio_name"] for s in samples],
+            "tgt_name": [s["tgt_name"] for s in samples],
+            "net_input": net_input,
+            "labels": labels,
+            "dec_target": collated_fbanks,
+            "dec_target_lengths": collated_fbanks_size,
+            "src_lengths": torch.LongTensor(audio_sizes),
+            "task_name": "s2s",
+            "ntokens": sum(audio_sizes),
+            "target": collated_fbanks,
+        }
+
+        return batch
+
+    def collater_audio(self, audios, audio_size):
+        collated_audios = audios[0].new_zeros(len(audios), audio_size)
+        padding_mask = (
+            torch.BoolTensor(collated_audios.shape).fill_(False)
+        )
+        for i, audio in enumerate(audios):
+            diff = len(audio) - audio_size
+            if diff == 0:
+                collated_audios[i] = audio
+            elif diff < 0:
+                collated_audios[i] = torch.cat([audio, audio.new_full((-diff,), 0.0)])
+                padding_mask[i, diff:] = True
+            else:
+                raise Exception("Diff should not be larger than 0")
+        return collated_audios, padding_mask
+
+
+    def num_tokens(self, index):
+        return self.wav_sizes[index]
+
+    def size(self, index):
+        return self.wav_sizes[index], self.tgt_sizes[index]
+
+    @property
+    def sizes(self):
+        return np.array(self.wav_sizes)
+
+    @property
+    def can_reuse_epoch_itr_across_epochs(self):
+        """No cache dataset if dataset is large-scale. Cache dataset for small dataset."""
+        return True
+
+    def ordered_indices(self):
+        if self.shuffle:
+            order = [np.random.permutation(len(self))]
+        else:
+            order = [np.arange(len(self))]
+
+        order.append(self.wav_sizes)
+        return np.lexsort(order)[::-1]
+
+    def postprocess(self, wav, cur_sample_rate):
+        if wav.dim() == 2:
+            wav = wav.mean(-1)
+        assert wav.dim() == 1, wav.dim()
+
+        if cur_sample_rate != self.sample_rate:
+            raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")
+
+        if self.normalize:
+            with torch.no_grad():
+                wav = F.layer_norm(wav, wav.shape)
+        return wav

artst/data/speech_to_text_dataset.py

@@ -0,0 +1,298 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import itertools
+import logging
+import os
+import mmap
+from typing import Any, List, Optional
+
+import numpy as np
+
+import torch
+torch.set_printoptions(profile="full")
+import torch.nn.functional as F
+from fairseq.data import data_utils, Dictionary
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+logger = logging.getLogger(__name__)
+
+
+def load_audio(manifest_path, max_keep, min_keep):
+    n_long, n_short = 0, 0
+    names, inds, sizes = [], [], []
+    with open(manifest_path) as f:
+        root = f.readline().strip()
+        for ind, line in enumerate(f):
+            items = line.strip().split("\t")
+            assert len(items) >= 2, line
+            sz = int(items[1])
+            if min_keep is not None and sz < min_keep:
+                n_short += 1
+            elif max_keep is not None and sz > max_keep:
+                n_long += 1
+            else:
+                names.append(items[0])
+                inds.append(ind)
+                sizes.append(sz)
+    tot = ind + 1
+    logger.info(
+        (
+            f"max_keep={max_keep}, min_keep={min_keep}, "
+            f"loaded {len(names)}, skipped {n_short} short and {n_long} long, "
+            f"longest-loaded={max(sizes)}, shortest-loaded={min(sizes)}"
+        )
+    )
+    return root, names, inds, tot, sizes
+
+
+def load_label(label_path, inds, tot):
+    with open(label_path) as f:
+        labels = [line.rstrip() for line in f]
+        assert (
+            len(labels) == tot
+        ), f"number of labels does not match ({len(labels)} != {tot})"
+        labels = [labels[i] for i in inds]
+    return labels
+
+
+def load_label_offset(label_path, inds, tot):
+    with open(label_path) as f:
+        # Hawau:
+        # changed line length reading as it's incorrect
+        code_lengths = [len(line.encode("utf-8")) for line in f] #original
+        # code_lengths = [len(line) for line in f] #fix
+        assert (
+            len(code_lengths) == tot
+        ), f"number of labels does not match ({len(code_lengths)} != {tot})"
+        offsets = list(itertools.accumulate([0] + code_lengths))
+        offsets = [(offsets[i], offsets[i + 1]) for i in inds]
+    return offsets
+
+
+class SpeechToTextDataset(FairseqDataset):
+    def __init__(
+        self,
+        manifest_path: str,
+        sample_rate: float,
+        label_paths: List[str],
+        label_processors: Optional[List[Any]] = None,
+        max_keep_sample_size: Optional[int] = None,
+        min_keep_sample_size: Optional[int] = None,
+        shuffle: bool = True,
+        normalize: bool = False,
+        store_labels: bool = True,
+        tgt_dict: Optional[Dictionary] = None,
+        tokenizer = None,
+    ):
+        self.audio_root, self.audio_names, inds, tot, self.wav_sizes = load_audio(
+            manifest_path, max_keep_sample_size, min_keep_sample_size
+        )
+       
+        self.sample_rate = sample_rate
+        self.shuffle = shuffle
+        self.tgt_dict = tgt_dict
+        self.tokenizer = tokenizer
+
+        self.num_labels = len(label_paths)
+        self.label_processors = label_processors
+        self.store_labels = store_labels
+
+        if store_labels:
+            self.label_list = [load_label(p, inds, tot) for p in label_paths]
+            logger.info(f"label_list: {self.label_list}")
+        else:
+            self.label_paths = label_paths
+            self.label_offsets_list = [
+                load_label_offset(p, inds, tot) for p in label_paths
+            ]
+            # logger.info(f"label_offsets_list: {self.label_offsets_list}")
+        assert label_processors is None or len(label_processors) == self.num_labels
+
+        self.normalize = normalize
+        logger.info(
+            f"normalize={normalize}"
+        )
+
+    def get_audio(self, index):
+        import soundfile as sf
+        # Hawau:
+        # logger.info(f"loaded_audio: {self.audio_names[index]}")
+        wav_path = os.path.join(self.audio_root, self.audio_names[index])
+        wav, cur_sample_rate = sf.read(wav_path)
+        wav = torch.from_numpy(wav).float()
+        wav = self.postprocess(wav, cur_sample_rate)
+        return wav
+
+    def get_label(self, index, label_idx):
+        if self.store_labels:
+            label = self.label_list[label_idx][index]
+        else:
+            # list slicing method
+            # with open(self.label_paths[label_idx]) as f:
+            #     offset_s, offset_e = self.label_offsets_list[label_idx][index]
+            #     # Hawau:
+            #     # f.seek(offset_s)
+            #     # label = f.read(offset_e - offset_s)
+            #     label = f.read()[offset_s : offset_e]
+            # Hawau:
+            # mmap method
+            with open(self.label_paths[label_idx], encoding='utf-8') as f:
+                with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+                    offset_s, offset_e = self.label_offsets_list[label_idx][index]
+                    label = mm[offset_s:offset_e].decode("utf-8")
+
+                
+        # Hawau:
+        # logger.info(f"loaded_label: {label}")
+        if self.tokenizer is not None:
+            label = self.tokenizer.encode(label)
+
+        if self.label_processors is not None:
+            label = self.label_processors[label_idx](label)
+        # logger.info(f"processed_label: {label}")
+        return label
+
+    def get_labels(self, index):
+        return [self.get_label(index, i) for i in range(self.num_labels)]
+
+    def __getitem__(self, index):
+        wav = self.get_audio(index)
+        labels = self.get_labels(index)
+        return {"id": index, "source": wav, "label_list": labels}
+
+    def __len__(self):
+        return len(self.wav_sizes)
+
+    def collater(self, samples):
+        samples = [s for s in samples if s["source"] is not None]
+        if len(samples) == 0:
+            return {}
+
+        audios = [s["source"] for s in samples]
+        audio_sizes = [len(s) for s in audios]
+
+        audio_size = max(audio_sizes)
+        collated_audios, padding_mask = self.collater_audio(
+            audios, audio_size
+        )
+
+        targets_by_label = [
+            [s["label_list"][i] for s in samples] for i in range(self.num_labels)
+        ]
+        targets_list, lengths_list, ntokens_list = self.collater_label(targets_by_label)
+
+        # Hawau:
+        # logger.info(f'targets_list: {targets_list}')
+    
+
+        decoder_label = [
+            torch.cat((targets_list[0][i, :lengths_list[0][i]], torch.tensor([self.tgt_dict.eos()])), 0).long()
+            for i in range(targets_list[0].size(0))
+        ]
+
+        decoder_target = data_utils.collate_tokens(
+            decoder_label,
+            self.tgt_dict.pad(),
+            self.tgt_dict.eos(),
+            left_pad=False,
+            move_eos_to_beginning=False,
+        )
+        decoder_target_lengths = torch.tensor(
+            [x.size(0) for x in decoder_label], dtype=torch.long
+        )
+        prev_output_tokens = data_utils.collate_tokens(
+            decoder_label,
+            self.tgt_dict.pad(),
+            self.tgt_dict.eos(),
+            left_pad=False,
+            move_eos_to_beginning=True,
+        )
+
+        net_input = {
+            "source": collated_audios, 
+            "padding_mask": padding_mask,
+            "prev_output_tokens": prev_output_tokens,
+            "task_name": "s2t",
+        }
+        batch = {
+            "id": torch.LongTensor([s["id"] for s in samples]),
+            "net_input": net_input,
+            "target": decoder_target,
+            "target_lengths": decoder_target_lengths,
+            "task_name": "s2t",
+            "ntokens": ntokens_list[0]
+        }
+
+        return batch
+
+    def collater_audio(self, audios, audio_size):
+        collated_audios = audios[0].new_zeros(len(audios), audio_size)
+        padding_mask = (
+            torch.BoolTensor(collated_audios.shape).fill_(False)
+        )
+        for i, audio in enumerate(audios):
+            diff = len(audio) - audio_size
+            if diff == 0:
+                collated_audios[i] = audio
+            elif diff < 0:
+                collated_audios[i] = torch.cat([audio, audio.new_full((-diff,), 0.0)])
+                padding_mask[i, diff:] = True
+            else:
+                raise Exception("Diff should not be larger than 0")
+        return collated_audios, padding_mask
+
+    def collater_seq_label(self, targets, pad):
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_label(self, targets_by_label):
+        targets_list, lengths_list, ntokens_list = [], [], []
+        itr = zip(targets_by_label, [self.tgt_dict.pad()])
+        
+        for targets, pad in itr:
+            # Hawau:
+            # logger.info(f'targets: {targets}')
+            targets, lengths, ntokens = self.collater_seq_label(targets, pad)
+            targets_list.append(targets)
+            lengths_list.append(lengths)
+            ntokens_list.append(ntokens)
+        return targets_list, lengths_list, ntokens_list
+
+    def num_tokens(self, index):
+        return self.size(index)
+
+    def size(self, index):
+        return self.wav_sizes[index]
+
+    @property
+    def sizes(self):
+        return np.array(self.wav_sizes)
+
+    def ordered_indices(self):
+        if self.shuffle:
+            order = [np.random.permutation(len(self))]
+        else:
+            order = [np.arange(len(self))]
+
+        order.append(self.wav_sizes)
+        return np.lexsort(order)[::-1]
+
+    def postprocess(self, wav, cur_sample_rate):
+        if wav.dim() == 2:
+            wav = wav.mean(-1)
+        assert wav.dim() == 1, wav.dim()
+
+        if cur_sample_rate != self.sample_rate:
+            raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")
+
+        if self.normalize:
+            with torch.no_grad():
+                wav = F.layer_norm(wav, wav.shape)
+        return wav

artst/data/text_dataset.py

@@ -0,0 +1,474 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import math
+
+import numpy as np
+import torch
+
+from fairseq.data import FairseqDataset, data_utils
+
+
+def collate(
+    samples,
+    pad_idx,
+    eos_idx,
+    vocab,
+    left_pad_source=False,
+    left_pad_target=False,
+    input_feeding=True,
+    pad_to_length=None,
+):
+    assert input_feeding
+    if len(samples) == 0:
+        return {}
+
+    def merge(key, left_pad, move_eos_to_beginning=False, pad_to_length=None):
+        return data_utils.collate_tokens(
+            [s[key] for s in samples],
+            pad_idx,
+            eos_idx=None,  # use eos_idx of each sample instead of vocab.eos()
+            left_pad=left_pad,
+            move_eos_to_beginning=move_eos_to_beginning,
+            pad_to_length=pad_to_length,
+        )
+
+    id = torch.LongTensor([s["id"] for s in samples])
+    src_tokens = merge(
+        "source",
+        left_pad=left_pad_source,
+        pad_to_length=pad_to_length["source"] if pad_to_length is not None else None,
+    )
+    # sort by descending source length
+    src_lengths = torch.LongTensor([s["source"].numel() for s in samples])
+    src_lengths, sort_order = src_lengths.sort(descending=True)
+    id = id.index_select(0, sort_order)
+    src_tokens = src_tokens.index_select(0, sort_order)
+
+    prev_output_tokens = None
+    target = None
+    if samples[0].get("target", None) is not None:
+        target = merge(
+            "target",
+            left_pad=left_pad_target,
+            pad_to_length=pad_to_length["target"]
+            if pad_to_length is not None
+            else None,
+        )
+        target = target.index_select(0, sort_order)
+        ntokens = sum(len(s["target"]) for s in samples)
+
+        if input_feeding:
+            # we create a shifted version of targets for feeding the
+            # previous output token(s) into the next decoder step
+            prev_output_tokens = merge(
+                "target",
+                left_pad=left_pad_target,
+                move_eos_to_beginning=True,
+                pad_to_length=pad_to_length["target"]
+                if pad_to_length is not None
+                else None,
+            )
+            prev_output_tokens = prev_output_tokens.index_select(0, sort_order)
+    else:
+        ntokens = sum(len(s["source"]) for s in samples)
+
+    batch = {
+        "id": id,
+        "ntokens": ntokens,
+        "net_input": {
+            "src_tokens": src_tokens,
+            "src_lengths": src_lengths,
+        },
+        "target": target,
+        "nsentences": samples[0]["source"].size(0),
+        "sort_order": sort_order,
+        "task_name": 'text_pretrain',
+    }
+    if prev_output_tokens is not None:
+        batch["net_input"]["prev_output_tokens"] = prev_output_tokens
+
+    return batch
+
+
+class TextPretrainDataset(FairseqDataset):
+    """
+    A wrapper around TokenBlockDataset for BART dataset.
+
+    Args:
+        dataset (TokenBlockDataset): dataset to wrap
+        sizes (List[int]): sentence lengths
+        vocab (~fairseq.data.Dictionary): vocabulary
+        mask_idx (int): dictionary index used for masked token
+        mask_whole_words: only mask whole words. This should be a byte mask
+            over vocab indices, indicating whether it is the beginning of a
+            word. We will extend any mask to encompass the whole word.
+        shuffle (bool, optional): shuffle the elements before batching.
+          Default: ``True``
+        seed: Seed for random number generator for reproducibility.
+        args: argparse arguments.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        sizes,
+        vocab,
+        mask_idx,
+        mask_whole_words,
+        shuffle,
+        seed,
+        args,
+        eos=None,
+        item_transform_func=None,
+        iid_noise_target=False,
+        uni_mask_idxs=None,
+    ):
+        self.dataset = dataset
+
+        self.sizes = sizes
+
+        self.vocab = vocab
+        self.shuffle = shuffle
+        self.seed = seed
+        if iid_noise_target:
+            assert isinstance(uni_mask_idxs, torch.Tensor), "if use iid_noise_target, the uni_mask_idxs must be a tensor which contain the mask indexs"
+        self.iid_noise_target = iid_noise_target
+        self.uni_mask_idxs = uni_mask_idxs
+        self.mask_idx = mask_idx
+        self.mask_whole_word = mask_whole_words
+        self.mask_ratio = args.mask
+        self.random_ratio = args.mask_random
+        self.insert_ratio = args.insert
+        self.rotate_ratio = args.rotate
+        self.permute_sentence_ratio = args.permute_sentences
+        self.eos = eos if eos is not None else vocab.eos()
+        self.item_transform_func = item_transform_func
+
+        if args.bpe != "gpt2":
+            self.full_stop_index = self.vocab.eos()
+        else:
+            assert args.bpe == "gpt2"
+            self.full_stop_index = self.vocab.index("13")
+
+        self.replace_length = args.replace_length
+        if self.replace_length not in [-1, 0, 1]:
+            raise ValueError(f"invalid arg: replace_length={self.replace_length}")
+        if args.mask_length not in ["subword", "word", "span-poisson"]:
+            raise ValueError(f"invalid arg: mask-length={args.mask_length}")
+        if args.mask_length == "subword" and args.replace_length not in [0, 1]:
+            raise ValueError(f"if using subwords, use replace-length=1 or 0")
+
+        self.mask_span_distribution = None
+        if args.mask_length == "span-poisson":
+            _lambda = args.poisson_lambda
+
+            lambda_to_the_k = 1
+            e_to_the_minus_lambda = math.exp(-_lambda)
+            k_factorial = 1
+            ps = []
+            for k in range(0, 128):
+                ps.append(e_to_the_minus_lambda * lambda_to_the_k / k_factorial)
+                lambda_to_the_k *= _lambda
+                k_factorial *= k + 1
+                if ps[-1] < 0.0000001:
+                    break
+            ps = torch.FloatTensor(ps)
+            self.mask_span_distribution = torch.distributions.Categorical(ps)
+
+        self.epoch = 0
+
+    @property
+    def can_reuse_epoch_itr_across_epochs(self):
+        return True  # only the noise changes, not item sizes
+
+    def set_epoch(self, epoch, **unused):
+        self.epoch = epoch
+
+    def __getitem__(self, index):
+        with data_utils.numpy_seed(self.seed, self.epoch, index):
+            tokens = self.dataset[index]
+            assert tokens[-1] == self.eos
+            source, target = tokens, tokens.clone()
+
+            if self.permute_sentence_ratio > 0.0:
+                source = self.permute_sentences(source, self.permute_sentence_ratio)
+
+            if self.mask_ratio > 0:
+                source, new_target = self.add_whole_word_mask(source, self.mask_ratio)
+                if new_target is not None:
+                    target = new_target
+
+            if self.insert_ratio > 0:
+                source = self.add_insertion_noise(source, self.insert_ratio)
+
+            if self.rotate_ratio > 0.0 and np.random.random() < self.rotate_ratio:
+                source = self.add_rolling_noise(source)
+        # there can additional changes to make:
+        if self.item_transform_func is not None:
+            source, target = self.item_transform_func(source, target)
+
+        assert (source >= 0).all()
+        assert (source[1:-1] >= 1).all()
+        assert (source <= len(self.vocab)).all()
+        assert source[0] == self.vocab.bos()
+        assert source[-1] == self.eos
+        return {
+            "id": index,
+            "source": source,
+            "target": target,
+        }
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def permute_sentences(self, source, p=1.0):
+        full_stops = source == self.full_stop_index
+        # Pretend it ends with a full stop so last span is a sentence
+        full_stops[-2] = 1
+
+        # Tokens that are full stops, where the previous token is not
+        sentence_ends = (full_stops[1:] * ~full_stops[:-1]).nonzero(as_tuple=False) + 2
+        result = source.clone()
+
+        num_sentences = sentence_ends.size(0)
+        num_to_permute = math.ceil((num_sentences * 2 * p) / 2.0)
+        substitutions = torch.randperm(num_sentences)[:num_to_permute]
+        ordering = torch.arange(0, num_sentences)
+        ordering[substitutions] = substitutions[torch.randperm(num_to_permute)]
+
+        # Ignore <bos> at start
+        index = 1
+        for i in ordering:
+            sentence = source[(sentence_ends[i - 1] if i > 0 else 1) : sentence_ends[i]]
+            result[index : index + sentence.size(0)] = sentence
+            index += sentence.size(0)
+        return result
+
+    def word_starts(self, source):
+        if self.mask_whole_word is not None:
+            is_word_start = self.mask_whole_word.gather(0, source)
+        else:
+            is_word_start = torch.ones(source.size())
+        is_word_start[0] = 0
+        is_word_start[-1] = 0
+        return is_word_start
+
+    def add_whole_word_mask(self, source, p):
+        source_ori = source.clone()
+        is_word_start = self.word_starts(source)
+        num_to_mask = int(math.ceil(is_word_start.float().sum() * p))
+        num_inserts = 0
+        if num_to_mask == 0:
+            return source
+
+        if self.mask_span_distribution is not None:
+            lengths = self.mask_span_distribution.sample(sample_shape=(num_to_mask,))
+
+            # Make sure we have enough to mask
+            cum_length = torch.cumsum(lengths, 0)
+            while cum_length[-1] < num_to_mask:
+                lengths = torch.cat(
+                    [
+                        lengths,
+                        self.mask_span_distribution.sample(sample_shape=(num_to_mask,)),
+                    ],
+                    dim=0,
+                )
+                cum_length = torch.cumsum(lengths, 0)
+
+            # Trim to masking budget
+            i = 0
+            while cum_length[i] < num_to_mask:
+                i += 1
+            lengths[i] = num_to_mask - (0 if i == 0 else cum_length[i - 1])
+            num_to_mask = i + 1
+            lengths = lengths[:num_to_mask]
+
+            # Handle 0-length mask (inserts) separately
+            lengths = lengths[lengths > 0]
+            num_inserts = num_to_mask - lengths.size(0)
+            num_to_mask -= num_inserts
+            if num_to_mask == 0:
+                return self.add_insertion_noise(source, num_inserts / source.size(0))
+
+            assert (lengths > 0).all()
+        else:
+            lengths = torch.ones((num_to_mask,)).long()
+        assert is_word_start[-1] == 0
+        word_starts = is_word_start.nonzero(as_tuple=False)
+        indices = word_starts[
+            torch.randperm(word_starts.size(0))[:num_to_mask]
+        ].squeeze(1)
+        mask_random = torch.FloatTensor(num_to_mask).uniform_() < self.random_ratio
+
+        source_length = source.size(0)
+        assert source_length - 1 not in indices
+        to_keep = torch.ones(source_length, dtype=torch.bool)
+        is_word_start[
+            -1
+        ] = 255  # acts as a long length, so spans don't go over the end of doc
+        if self.replace_length == 0:
+            to_keep[indices] = 0
+        else:
+            # keep index, but replace it with [MASK]
+            source[indices] = self.mask_idx
+            source[indices[mask_random]] = torch.randint(
+                1, len(self.vocab), size=(mask_random.sum(),)
+            )
+
+        if self.mask_span_distribution is not None:
+            assert len(lengths.size()) == 1
+            assert lengths.size() == indices.size()
+            lengths -= 1
+            while indices.size(0) > 0:
+                assert lengths.size() == indices.size()
+                lengths -= is_word_start[indices + 1].long()
+                uncompleted = lengths >= 0
+                indices = indices[uncompleted] + 1
+                mask_random = mask_random[uncompleted]
+                lengths = lengths[uncompleted]
+                if self.replace_length != -1:
+                    # delete token
+                    to_keep[indices] = 0
+                else:
+                    # keep index, but replace it with [MASK]
+                    source[indices] = self.mask_idx
+                    source[indices[mask_random]] = torch.randint(
+                        1, len(self.vocab), size=(mask_random.sum(),)
+                    )
+        else:
+            # A bit faster when all lengths are 1
+            while indices.size(0) > 0:
+                uncompleted = is_word_start[indices + 1] == 0
+                indices = indices[uncompleted] + 1
+                mask_random = mask_random[uncompleted]
+                if self.replace_length != -1:
+                    # delete token
+                    to_keep[indices] = 0
+                else:
+                    # keep index, but replace it with [MASK]
+                    source[indices] = self.mask_idx
+                    source[indices[mask_random]] = torch.randint(
+                        1, len(self.vocab), size=(mask_random.sum(),)
+                    )
+
+                assert source_length - 1 not in indices
+
+        if not self.iid_noise_target:
+            source = source[to_keep]
+            target = None
+        else:
+            ## Prepare source
+            source_mask_idx = (source == self.mask_idx).nonzero().view(-1)
+            source[source_mask_idx] = self.uni_mask_idxs[:source_mask_idx.size(0)]
+            source = source[to_keep]
+
+            ## Prepare target
+            to_keep[source_mask_idx] = 0
+
+            # source_mask_idx: from [a, b, c, ...] to [a, b + 1, c + 2, ...]
+            source_mask_idx = source_mask_idx + torch.arange(source_mask_idx.size(0))
+            # target: source_length + mask_length
+            target = source_ori.new_zeros(source_mask_idx.size(0) + source_ori.size(0))
+            # target: [0, 0, 0, X, 0, 0, Y, ....]
+            target[source_mask_idx] = self.uni_mask_idxs[:source_mask_idx.size(0)]
+
+            target_to_keep = to_keep.new_zeros(source_mask_idx.size(0) + source_ori.size(0))
+
+            # Copy original value to target and target_to_keep
+            target_to_keep[target == 0] = to_keep
+            target_to_keep[-1] = 0
+            target[target == 0] = source_ori
+
+            target = target[~target_to_keep]
+
+        if num_inserts > 0:
+            source = self.add_insertion_noise(source, num_inserts / source.size(0))
+
+        return source, target
+
+    def add_permuted_noise(self, tokens, p):
+        num_words = len(tokens)
+        num_to_permute = math.ceil(((num_words * 2) * p) / 2.0)
+        substitutions = torch.randperm(num_words - 2)[:num_to_permute] + 1
+        tokens[substitutions] = tokens[substitutions[torch.randperm(num_to_permute)]]
+        return tokens
+
+    def add_rolling_noise(self, tokens):
+        offset = np.random.randint(1, max(1, tokens.size(-1) - 1) + 1)
+        tokens = torch.cat(
+            (tokens[0:1], tokens[offset:-1], tokens[1:offset], tokens[-1:]),
+            dim=0,
+        )
+        return tokens
+
+    def add_insertion_noise(self, tokens, p):
+        if p == 0.0:
+            return tokens
+
+        num_tokens = len(tokens)
+        n = int(math.ceil(num_tokens * p))
+
+        noise_indices = torch.randperm(num_tokens + n - 2)[:n] + 1
+        noise_mask = torch.zeros(size=(num_tokens + n,), dtype=torch.bool)
+        noise_mask[noise_indices] = 1
+        result = torch.LongTensor(n + len(tokens)).fill_(-1)
+
+        num_random = int(math.ceil(n * self.random_ratio))
+        result[noise_indices[num_random:]] = self.mask_idx
+        result[noise_indices[:num_random]] = torch.randint(
+            low=1, high=len(self.vocab), size=(num_random,)
+        )
+
+        result[~noise_mask] = tokens
+
+        assert (result >= 0).all()
+        return result
+
+    def collater(self, samples, pad_to_length=None):
+        """Merge a list of samples to form a mini-batch.
+        Args:
+            samples (List[dict]): samples to collate
+        Returns:
+            dict: a mini-batch of data
+        """
+        return collate(
+            samples, self.vocab.pad(), self.eos, self.vocab, pad_to_length=pad_to_length
+        )
+
+    def num_tokens(self, index):
+        """Return the number of tokens in a sample. This value is used to
+        enforce ``--max-tokens`` during batching."""
+        return self.sizes[index]
+
+    def size(self, index):
+        """Return an example's size as a float or tuple. This value is used when
+        filtering a dataset with ``--max-positions``."""
+        return self.sizes[index]
+
+    def ordered_indices(self):
+        """Return an ordered list of indices. Batches will be constructed based
+        on this order."""
+        if self.shuffle:
+            indices = np.random.permutation(len(self))
+        else:
+            indices = np.arange(len(self))
+        return indices[np.argsort(self.sizes[indices], kind="mergesort")]
+
+    def prefetch(self, indices):
+        self.src.prefetch(indices)
+        self.tgt.prefetch(indices)
+
+    @property
+    def supports_prefetch(self):
+        return (
+            hasattr(self.src, "supports_prefetch")
+            and self.src.supports_prefetch
+            and hasattr(self.tgt, "supports_prefetch")
+            and self.tgt.supports_prefetch
+        )

artst/data/text_to_speech_dataset.py

@@ -0,0 +1,344 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import itertools
+import logging
+import os
+from typing import Any, List, Optional
+import mmap
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+import librosa
+from fairseq.data.audio.speech_to_text_dataset import get_features_or_waveform
+from fairseq.data import data_utils, Dictionary
+from fairseq.data.fairseq_dataset import FairseqDataset
+
+
+logger = logging.getLogger(__name__)
+
+def _collate_frames(
+    frames: List[torch.Tensor], is_audio_input: bool = False
+):
+    """
+    Convert a list of 2D frames into a padded 3D tensor
+    Args:
+        frames (list): list of 2D frames of size L[i]*f_dim. Where L[i] is
+            length of i-th frame and f_dim is static dimension of features
+    Returns:
+        3D tensor of size len(frames)*len_max*f_dim where len_max is max of L[i]
+    """
+    max_len = max(frame.size(0) for frame in frames)
+    if is_audio_input:
+        out = frames[0].new_zeros((len(frames), max_len))
+    else:
+        out = frames[0].new_zeros((len(frames), max_len, frames[0].size(1)))
+    for i, v in enumerate(frames):
+        out[i, : v.size(0)] = v
+    return out
+
+def load_audio(manifest_path, max_keep, min_keep):
+    n_long, n_short = 0, 0
+    names, inds, sizes, spk_embeds = [], [], [], []
+    with open(manifest_path) as f:
+        root = f.readline().strip()
+        for ind, line in enumerate(f):
+            items = line.strip().split("\t")
+            assert len(items) == 3, line
+            sz = int(items[1])
+            if min_keep is not None and sz < min_keep:
+                n_short += 1
+            elif max_keep is not None and sz > max_keep:
+                n_long += 1
+            else:
+                names.append(items[0])
+                spk_embeds.append(items[2])
+                inds.append(ind)
+                sizes.append(sz)
+    tot = ind + 1
+    logger.info(
+        (
+            f"max_keep={max_keep}, min_keep={min_keep}, "
+            f"loaded {len(names)}, skipped {n_short} short and {n_long} long, "
+            f"longest-loaded={max(sizes)}, shortest-loaded={min(sizes)}"
+        )
+    )
+    return root, names, inds, tot, sizes, spk_embeds
+
+
+def load_label(label_path, inds, tot):
+    with open(label_path) as f:
+        labels = [line.rstrip() for line in f]
+        assert (
+            len(labels) == tot
+        ), f"number of labels does not match ({len(labels)} != {tot})"
+        labels = [labels[i] for i in inds]
+    return labels
+
+
+def load_label_offset(label_path, inds, tot):
+    with open(label_path, encoding='utf-8') as f:
+        code_lengths = [len(line.encode("utf-8")) for line in f] #changed as in speech_to_text_dataset.py
+        assert (
+            len(code_lengths) == tot
+        ), f"number of labels does not match ({len(code_lengths)} != {tot})"
+        offsets = list(itertools.accumulate([0] + code_lengths))
+        offsets = [(offsets[i], offsets[i + 1]) for i in inds]
+    return offsets
+
+
+def logmelfilterbank(
+    audio,
+    sampling_rate,
+    fft_size=1024,
+    hop_size=256,
+    win_length=None,
+    window="hann",
+    num_mels=80,
+    fmin=80,
+    fmax=7600,
+    eps=1e-10,
+):
+    """Compute log-Mel filterbank feature. 
+    (https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/bin/preprocess.py)
+
+    Args:
+        audio (ndarray): Audio signal (T,).
+        sampling_rate (int): Sampling rate.
+        fft_size (int): FFT size.
+        hop_size (int): Hop size.
+        win_length (int): Window length. If set to None, it will be the same as fft_size.
+        window (str): Window function type.
+        num_mels (int): Number of mel basis.
+        fmin (int): Minimum frequency in mel basis calculation.
+        fmax (int): Maximum frequency in mel basis calculation.
+        eps (float): Epsilon value to avoid inf in log calculation.
+
+    Returns:
+        ndarray: Log Mel filterbank feature (#frames, num_mels).
+
+    """
+    # get amplitude spectrogram
+    x_stft = librosa.stft(audio, n_fft=fft_size, hop_length=hop_size,
+                          win_length=win_length, window=window, pad_mode="reflect")
+    spc = np.abs(x_stft).T  # (#frames, #bins)
+
+    # get mel basis
+    fmin = 0 if fmin is None else fmin
+    fmax = sampling_rate / 2 if fmax is None else fmax
+    mel_basis = librosa.filters.mel(sr=sampling_rate, n_fft=fft_size, n_mels=num_mels, fmin=fmin, fmax=fmax)
+
+    return np.log10(np.maximum(eps, np.dot(spc, mel_basis.T)))
+
+
+
+class TextToSpeechDataset(FairseqDataset):
+    def __init__(
+        self,
+        manifest_path: str,
+        sample_rate: float,
+        label_paths: List[str],
+        label_processors: Optional[List[Any]] = None,
+        max_keep_sample_size: Optional[int] = None,
+        min_keep_sample_size: Optional[int] = None,
+        shuffle: bool = True,
+        normalize: bool = False,
+        store_labels: bool = True,
+        src_dict: Optional[Dictionary] = None,
+        tokenizer = None,
+        reduction_factor: int = 1,
+        inference: bool = False,
+    ):
+        
+        self.audio_root, self.audio_names, inds, tot, self.wav_sizes, self.spk_embeds = load_audio(
+            manifest_path, max_keep_sample_size, min_keep_sample_size
+        )
+        self.inference = inference
+        
+        self.sample_rate = sample_rate
+        self.shuffle = shuffle
+        self.src_dict = src_dict
+        self.tokenizer = tokenizer
+
+        self.num_labels = len(label_paths)
+        self.label_processors = label_processors
+        self.store_labels = store_labels
+        if store_labels:
+            self.label_list = [load_label(p, inds, tot) for p in label_paths]
+        else:
+            self.label_paths = label_paths
+            self.label_offsets_list = [
+                load_label_offset(p, inds, tot) for p in label_paths
+            ]
+        assert label_processors is None or len(label_processors) == self.num_labels
+
+        self.normalize = normalize
+        self.reduction_factor = reduction_factor
+        logger.info(
+            f"reduction_factor={reduction_factor}, normalize={normalize}"
+        )
+
+    def get_audio(self, index):
+        import soundfile as sf
+
+        wav_path = os.path.join(self.audio_root, self.audio_names[index])
+        wav, cur_sample_rate = sf.read(wav_path)
+        wav = torch.from_numpy(wav).float()
+        fbank = logmelfilterbank(
+            wav.view(-1).cpu().numpy(), 16000
+        )
+        fbank = torch.from_numpy(fbank).float()
+        wav = self.postprocess(wav, cur_sample_rate)
+        return wav, fbank
+
+    def get_label(self, index, label_idx):
+        if self.store_labels:
+            label = self.label_list[label_idx][index]
+        else:
+            # with open(self.label_paths[label_idx]) as f:
+            #     offset_s, offset_e = self.label_offsets_list[label_idx][index]
+            #     f.seek(offset_s)
+            #     label = f.read(offset_e - offset_s)
+            
+            # Hawau:
+            # mmap method
+            with open(self.label_paths[label_idx], encoding='utf-8') as f:
+                with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+                    offset_s, offset_e = self.label_offsets_list[label_idx][index]
+                    label = mm[offset_s:offset_e].decode("utf-8")
+
+
+        if self.tokenizer is not None:
+            label = self.tokenizer.encode(label)
+
+        if self.label_processors is not None:
+            label = self.label_processors[label_idx](label)
+        return label
+
+    def get_labels(self, index):
+        return [self.get_label(index, i) for i in range(self.num_labels)]
+
+    def __getitem__(self, index):
+        wav, fbank = self.get_audio(index)            
+        labels = self.get_labels(index)
+        spkembs = get_features_or_waveform(
+            os.path.join(self.audio_root, self.spk_embeds[index])
+        )
+        spkembs = torch.from_numpy(spkembs).float()
+        
+        return {"id": index, "source": labels, "target": fbank, "spkembs": spkembs, "audio_name": self.audio_names[index]}
+
+    
+    def __len__(self):
+        return len(self.wav_sizes)
+
+    def collater(self, samples):
+        samples = [s for s in samples if s["source"] is not None]
+        if len(samples) == 0:
+            return {}
+
+        fbanks = [s["target"] for s in samples]
+        fbank_sizes = [len(s) for s in fbanks]
+
+        collated_fbanks = _collate_frames(fbanks)
+        collated_fbanks_size = torch.tensor(fbank_sizes, dtype=torch.long)
+
+        # thin out frames for reduction factor (B, Lmax, odim) ->  (B, Lmax//r, odim)
+        if self.reduction_factor > 1:
+            collated_fbanks_in = collated_fbanks[:, self.reduction_factor - 1 :: self.reduction_factor]
+            collated_fbanks_size_in = collated_fbanks_size.new([torch.div(olen, self.reduction_factor, rounding_mode='floor') for olen in collated_fbanks_size])
+        else:
+            collated_fbanks_in, collated_fbanks_size_in = collated_fbanks, collated_fbanks_size
+
+        prev_output_tokens = torch.cat(
+            [collated_fbanks_in.new_zeros((collated_fbanks_in.shape[0], 1, collated_fbanks_in.shape[2])), collated_fbanks_in[:, :-1]], dim=1
+        )
+
+        # make labels for stop prediction
+        labels = collated_fbanks.new_zeros(collated_fbanks.size(0), collated_fbanks.size(1))
+        for i, l in enumerate(fbank_sizes):
+            labels[i, l - 1 :] = 1.0
+
+        spkembs = _collate_frames([s["spkembs"] for s in samples], is_audio_input=True)
+
+        sources_by_label = [
+            [s["source"][i] for s in samples] for i in range(self.num_labels)
+        ]
+        sources_list, lengths_list, ntokens_list = self.collater_label(sources_by_label)
+
+        net_input = {
+            "src_tokens": sources_list[0], 
+            "src_lengths": lengths_list[0],
+            "prev_output_tokens": prev_output_tokens,
+            "tgt_lengths": collated_fbanks_size_in,
+            "spkembs": spkembs,
+            "task_name": "t2s",
+        }
+        batch = {
+            "id": torch.LongTensor([s["id"] for s in samples]),
+            "name": [s["audio_name"] for s in samples],
+            "net_input": net_input,
+            "labels": labels,
+            "dec_target": collated_fbanks,
+            "dec_target_lengths": collated_fbanks_size,
+            "src_lengths": lengths_list[0],
+            "task_name": "t2s",
+            "ntokens": ntokens_list[0],
+            "target": collated_fbanks,
+        }
+
+        return batch
+
+    def collater_seq_label(self, targets, pad):
+        lengths = torch.LongTensor([len(t) for t in targets])
+        ntokens = lengths.sum().item()
+        targets = data_utils.collate_tokens(targets, pad_idx=pad, left_pad=False)
+        return targets, lengths, ntokens
+
+    def collater_label(self, targets_by_label):
+        targets_list, lengths_list, ntokens_list = [], [], []
+        itr = zip(targets_by_label, [self.src_dict.pad()])
+        for targets, pad in itr:
+            targets, lengths, ntokens = self.collater_seq_label(targets, pad)
+            targets_list.append(targets)
+            lengths_list.append(lengths)
+            ntokens_list.append(ntokens)
+        return targets_list, lengths_list, ntokens_list
+
+    def num_tokens(self, index):
+        return self.size(index)
+
+    def size(self, index):
+        return self.wav_sizes[index]
+
+    @property
+    def sizes(self):
+        return np.array(self.wav_sizes)
+
+    def ordered_indices(self):
+        if self.shuffle:
+            order = [np.random.permutation(len(self))]
+        else:
+            order = [np.arange(len(self))]
+
+        order.append(self.wav_sizes)
+        return np.lexsort(order)[::-1]
+
+    def postprocess(self, wav, cur_sample_rate):
+        if wav.dim() == 2:
+            wav = wav.mean(-1)
+        assert wav.dim() == 1, wav.dim()
+
+        if cur_sample_rate != self.sample_rate:
+            raise Exception(f"sr {cur_sample_rate} != {self.sample_rate}")
+
+        if self.normalize:
+            with torch.no_grad():
+                wav = F.layer_norm(wav, wav.shape)
+        return wav

artst/models/__init__.py

@@ -0,0 +1,2 @@
+from .artst import *  # noqa
+from .t5_transformer_lm import *  # noqa

artst/models/artst.py

@@ -0,0 +1,1448 @@
+# --------------------------------------------------------
+# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
+# Github source: https://github.com/mbzuai-nlp/ArTST
+# Based on speecht5, fairseq and espnet code bases
+# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
+# --------------------------------------------------------
+
+import logging
+from ast import literal_eval
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from fairseq import utils
+from fairseq.models import (
+    FairseqEncoderDecoderModel,
+    FairseqIncrementalDecoder,
+    register_model,
+    register_model_architecture,
+)
+from .modules.text_encoder_prenet import TextEncoderPrenet
+from .modules.text_decoder_prenet import TextDecoderPrenet
+from .modules.text_decoder_postnet import TextDecoderPostnet
+from .modules.speech_encoder_prenet import SpeechEncoderPrenet
+from .modules.speech_encoder_postnet import SpeechEncoderPostnet
+from .modules.speech_decoder_prenet import SpeechDecoderPrenet
+from .modules.speech_decoder_postnet import SpeechDecoderPostnet
+from .modules.speaker_decoder_postnet import SpeakerDecoderPostnet
+from .modules.encoder import TransformerEncoder
+from .modules.decoder import TransformerDecoder
+from fairseq.modules.transformer_sentence_encoder import init_bert_params
+from fairseq.models.transformer import Embedding
+from fairseq.modules import (
+    GumbelVectorQuantizer,
+)
+from torch import Tensor
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_TEXT_POSITIONS = 450
+DEFAULT_MAX_SPEECH_POSITIONS = 4000
+
+
+@register_model("artst_transformer")
+class ArTSTTransformerModel(FairseqEncoderDecoderModel):
+    """Adapted Transformer model (https://arxiv.org/abs/1706.03762) for
+    speech-to-text tasks. The Transformer encoder/decoder remains the same.
+    A trainable input subsampler is prepended to the Transformer encoder to
+    project inputs into the encoder dimension as well as downsample input
+    sequence for computational efficiency."""
+
+    def __init__(
+            self, 
+            args,
+            encoder, decoder,
+            text_encoder_prenet, speech_encoder_prenet,
+            text_decoder_prenet, speech_decoder_prenet,
+            text_decoder_postnet, speech_decoder_postnet,
+            speaker_decoder_postnet, speech_encoder_postnet, 
+        ):
+        super().__init__(encoder, decoder)
+
+        self.encoder = encoder
+        self.decoder = decoder
+
+        self.text_encoder_prenet = text_encoder_prenet
+        self.speech_encoder_prenet = speech_encoder_prenet
+
+        self.text_decoder_prenet = text_decoder_prenet
+        self.speech_decoder_prenet = speech_decoder_prenet
+
+        self.text_decoder_postnet = text_decoder_postnet
+        self.speech_decoder_postnet = speech_decoder_postnet
+        self.speaker_decoder_postnet = speaker_decoder_postnet
+
+        self.hubert_layer = speech_encoder_postnet
+
+        self.reduction_factor = args.reduction_factor
+        self.spk_embed_dim = args.spk_embed_dim
+
+        # define projection layer
+        self.spk_embed_integration_type = args.spk_embed_integration_type
+        if self.spk_embed_dim is not None and self.spk_embed_integration_type != 'pre':
+            if self.spk_embed_integration_type == "add":
+                self.projection = torch.nn.Linear(self.spk_embed_dim, args.decoder_embed_dim)
+            else:
+                self.projection = torch.nn.Linear(
+                    args.decoder_embed_dim + self.spk_embed_dim, args.decoder_embed_dim
+                )
+
+        # Hawau: here we can add language embedding integration
+
+        self.use_codebook = args.use_codebook
+        self.codebook_prob = getattr(args, "codebook_prob", 0.5) # args.codebook_prob
+        if self.use_codebook:
+            vq_dim = args.latent_dim if args.latent_dim > 0 else args.encoder_embed_dim
+            self.quantizer = GumbelVectorQuantizer(
+                dim=args.encoder_embed_dim,
+                num_vars=args.latent_vars,
+                temp=args.latent_temp,
+                groups=args.latent_groups,
+                combine_groups=False,
+                vq_dim=vq_dim,
+                time_first=True,
+                weight_proj_depth=args.quantizer_depth,
+                weight_proj_factor=args.quantizer_factor,
+            )
+
+        self.num_updates = 0
+
+        # # Follow BERT's random weight initialization (for BART)
+        if args.bert_init:
+            self.apply(init_bert_params)
+        self.args = args
+        self.prune_modules(args.modules_filter)
+
+    @staticmethod
+    def add_args(parser):
+        """Add model-specific arguments to the parser."""
+        # Transformer
+        parser.add_argument(
+            "--activation-fn",
+            type=str,
+            choices=utils.get_available_activation_fns(),
+            help="activation function to use",
+        )
+        parser.add_argument(
+            "--dropout", type=float, metavar="D", help="dropout probability"
+        )
+        parser.add_argument(
+            "--attention-dropout",
+            type=float,
+            metavar="D",
+            help="dropout probability for attention weights",
+        )
+        parser.add_argument(
+            "--activation-dropout",
+            "--relu-dropout",
+            type=float,
+            metavar="D",
+            help="dropout probability after activation in FFN.",
+        )
+        parser.add_argument(
+            "--encoder-embed-dim",
+            type=int,
+            metavar="N",
+            help="encoder embedding dimension",
+        )
+        parser.add_argument(
+            "--encoder-ffn-embed-dim",
+            type=int,
+            metavar="N",
+            help="encoder embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--encoder-layers", type=int, metavar="N", help="num encoder layers"
+        )
+        parser.add_argument(
+            "--encoder-attention-heads",
+            type=int,
+            metavar="N",
+            help="num encoder attention heads",
+        )
+        parser.add_argument(
+            "--encoder-normalize-before",
+            action="store_true",
+            help="apply layernorm before each encoder block",
+        )
+        parser.add_argument(
+            "--decoder-normalize-before",
+            action="store_true",
+            help="apply layernorm before each decoder block",
+        )
+        parser.add_argument(
+            "--decoder-embed-dim",
+            type=int,
+            metavar="N",
+            help="decoder embedding dimension",
+        )
+        parser.add_argument(
+            "--decoder-ffn-embed-dim",
+            type=int,
+            metavar="N",
+            help="decoder embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--decoder-layers", type=int, metavar="N", help="num decoder layers"
+        )
+        parser.add_argument(
+            "--decoder-attention-heads",
+            type=int,
+            metavar="N",
+            help="num decoder attention heads",
+        )
+        parser.add_argument(
+            "--reduction-factor",
+            type=int,
+            help="reduction factor for decoder",
+        )
+        parser.add_argument(
+            "--spk-embed-dim",
+            type=int,
+            help="speaker embedding dimension",
+        )
+        parser.add_argument(
+            "--layernorm-embedding",
+            action="store_true",
+            help="add layernorm to embedding",
+        )
+        parser.add_argument(
+            "--load-pretrained-encoder-from",
+            type=str,
+            metavar="STR",
+            help="model to take encoder weights from (for initialization)",
+        )
+        parser.add_argument(
+            '--freeze-encoder-updates',
+            type=int,
+            help='number of steps to freeze encoder before finetune'
+        )
+        parser.add_argument(
+            '--freeze-decoder-updates',
+            type=int,
+            help='number of steps to freeze decoder before finetune'
+        )
+        parser.add_argument(
+            '--no-freeze-encoder-layer',
+            type=str,
+            help='which encoder layer not freeze during finetune'
+        )
+        parser.add_argument(
+            "--share-input-output-embed",
+            action="store_true",
+            help="share decoder input and output embeddings",
+        )
+        parser.add_argument(
+            "--share-ctc-embed",
+            action="store_true",
+            help="share ctc embed and decoder embed",
+        )
+        parser.add_argument(
+            "--encoder-sliding-window-attn",
+            default=None,
+            type=int,
+            help="If not None but a even number, set sliding window attention to encoder's attn_mask, e.g., 4, 10, and 20",
+        )
+        
+        # Convolutional subsampler
+        parser.add_argument(
+            "--encoder-speech-prenet",
+            default="conv",
+            type=str,
+            choices=["conv", "linear"],
+            help="The type of encoder speech prenet, e.g., conv or linear."
+        )
+        parser.add_argument(
+            "--conv-kernel-sizes",
+            default="5,5",
+            type=str,
+            help="The layer of convolution of encoder speech prenet."
+        )
+        parser.add_argument(
+            "--conv-channels",
+            default=1024,
+            type=int,
+            help="The channels of encoder speech prenet."
+        )
+        parser.add_argument(
+            "--subsample-stride",
+            default="2,2",
+            type=str,
+            help="The subsample stride for conv1dsubsample."
+        )
+        parser.add_argument(
+            "--spk-embed-integration-type",
+            type=str,
+            choices=["pre", "add"],
+            help="speaker embedding integration type"
+        )
+        parser.add_argument(
+            "--dprenet-dropout-rate",
+            default=0.5,
+            type=float,
+            help="The dropout rate of decoder speech prenet."
+        )
+        
+        ## SE
+        parser.add_argument(
+            "--se-predict",
+            default=None,
+            choices=["masking", "target", "delta"],
+            help="If set, source speech inputs decoder to predict the masking/target/delta of corresponding inputs."
+               + "masking is [0, 1], target is predicted output, delta is difference between inputs and outputs",
+        )
+        parser.add_argument(
+            "--se-decoder-input",
+            type=str,
+            default="previous_target",
+            choices=["previous_target", "source"],
+        )
+        
+        ## SID
+        parser.add_argument(
+            "--modules-filter",
+            default=None,
+            type=str,
+            help="Remove unused modules for, e.g., SID.",
+        )
+        parser.add_argument(
+            "--sid-pad-prenet",
+            action="store_true",
+            help="If set, the size of text dictionary is as small as for <pad> token.",
+        )
+        parser.add_argument(
+            "--encoder-attn-branch",
+            type=str,
+            default="identity,full",
+            help="encoder attention branch sliding window, e.g., 'identity,0,2,4,full'",
+        )
+        parser.add_argument(
+            "--encoder-block-branch",
+            type=str,
+            help="average the output of encoder, e.g., '4,5,6'",
+        )
+        parser.add_argument(
+            "--sid-encoder-cls",
+            default=None,
+            choices=["encoder"],
+            help="If set, add cls vector to the encoder input, e.g., constant vector.",
+        )
+        parser.add_argument(
+            "--sid-shuffle-encoder-input",
+            action="store_true",
+            help="If set, shuffle encoder input in time.",
+        )
+        parser.add_argument(
+            "--sid-decoder-speaker",
+            action="store_true",
+            help="If set, apply speaker decoder as transformer decoder.",
+        )
+        parser.add_argument(
+            "--sid-decoder-attn-dim",
+            default=128,
+            type=int,
+            help="Attention dimension in attensive statistics pooling of speaker decoder.",
+        )
+        parser.add_argument(
+            "--sid-t5-postnet",
+            action="store_true",
+            help="If set, apply TextDecoderPostnet as speaker classification.",
+        )
+        parser.add_argument(
+            "--sid-embed-dim",
+            default=128,
+            type=int,
+            help="Embedding dimension in speaker postnet for speaker identification if embed postnet.",
+        )
+        parser.add_argument(
+            "--sid-pooling-layer",
+            default="decoder",
+            type=str,
+            choices=["decoder-las", "decoder", "encoder", "encoder-cls", "encoder-speaker"],
+            help="The output of decoder or encoder uses as SID pooling layer over temporal dimension.",
+        )
+        parser.add_argument(
+            "--sid-no-pooling-bn",
+            action="store_true",
+            help="If set, not attention batchnorm.",
+        )
+        parser.add_argument(
+            "--sid-no-embed-postnet",
+            action="store_true",
+            help="If set, no layer between decoder output and classification layer.",
+        )
+        parser.add_argument(
+            "--sid-normalize-postnet",
+            action="store_true",
+            help="If set, normalize input and weight in postnet/classifier.",
+        )
+        parser.add_argument(
+            "--sid-softmax-type",
+            default="softmax",
+            choices=["softmax", "amsoftmax", "aamsoftmax"],
+            help="If using amsoftmax or aamsoftmax, the target should be given.",
+        )
+        parser.add_argument(
+            "--softmax-scale",
+            default=1.0,
+            type=float,
+            help="Scale for AMSoftmax or AAMSoftmax.",
+        )
+        parser.add_argument(
+            "--softmax-margin",
+            default=0.0,
+            type=float,
+            help="Margin for AMSoftmax or AAMSoftmax.",
+        )
+        parser.add_argument(
+            "--softmax-easy-margin",
+            action="store_true",
+            help="Enable easy margin for AAMSoftmax.",
+        )
+        parser.add_argument(
+            "--encoder-layerdrop",
+            type=float,
+            metavar="D",
+            help="LayerDrop probability for encoder",
+        )
+        parser.add_argument(
+            "--decoder-layerdrop",
+            type=float,
+            metavar="D",
+            help="LayerDrop probability for decoder",
+        )
+        
+        ## Hubert
+        parser.add_argument(
+            '--feature-grad-mult',
+            type=float,
+            help='multiply feature extractor var grads by this'
+        )
+        parser.add_argument(
+            '--logit-temp',
+            type=float,
+            help='temperature to divide logits by'
+        )
+        parser.add_argument(
+            '--final-dim',
+            type=int,
+            help="project final representations and targets to this many "
+            "dimensions. set to encoder_embed_dim is <= 0"
+        )
+        
+        # mask
+        parser.add_argument(
+            '--hubert-mask-length',
+            type=int,
+            help='mask length'
+        )
+        parser.add_argument(
+            '--mask-prob',
+            type=float,
+            help='probability of replacing a token with mask'
+        )
+        parser.add_argument(
+            "--mask-selection",
+            choices=["static", "uniform", "normal", "poisson"],
+            help="how to choose mask length",
+        )
+        parser.add_argument(
+            '--mask-other',
+            type=float,
+            help="secondary mask argument "
+            "(used for more complex distributions), "
+            "see help in compute_mask_indices"
+        )
+        parser.add_argument(
+            '--mask-min-space',
+            type=int,
+            help='min space between spans (if no overlap is enabled)'
+        )
+        
+        # channel masking
+        parser.add_argument(
+            '--mask-channel-length',
+            type=int,
+            help='length of the mask for features (channels)'
+        )
+        parser.add_argument(
+            '--mask-channel-prob',
+            type=float,
+            help="probability of replacing a feature with 0"
+        )
+        parser.add_argument(
+            "--mask-channel-selection",
+            choices=["static", "uniform", "normal", "poisson"],
+            help="how to choose mask length for channel masking",
+        )
+        parser.add_argument(
+            '--mask-channel-other',
+            type=float,
+            help="secondary mask argument "
+            "(used for more complex distributions), "
+            "see help in compute_mask_indices"
+        )
+        parser.add_argument(
+            '--mask-channel-min-space',
+            type=int,
+            help='min space between spans (if no overlap is enabled)'
+        )
+        
+        # abs positional embeddings
+        parser.add_argument(
+            '--conv-pos',
+            type=int,
+            help='number of filters for convolutional positional embeddings'
+        )
+        parser.add_argument(
+            '--conv-pos-groups',
+            type=int,
+            help='number of groups for convolutional positional embedding'
+        )
+        
+        # codebook related
+        parser.add_argument(
+            "--use-codebook",
+            action="store_true",
+            help="whether to use codebook",
+        )
+        parser.add_argument(
+            "--codebook-prob",
+            type=float,
+            help="probability to use codebook",
+        )
+        parser.add_argument(
+            "--latent-vars",
+            type=int,
+            help="number of latent variables V in each group of the codebook",
+        )
+        parser.add_argument(
+            "--latent-groups",
+            type=int,
+            help="number of groups G of latent variables in the codebook",
+        )
+        parser.add_argument(
+            "--latent-dim",
+            type=int,
+            help="if > 0, uses this dimensionality for latent variables. "
+            "otherwise uses final_dim / latent_groups",
+        )
+        parser.add_argument(
+            "--latent-temp",
+            type=literal_eval,
+            help="temperature for latent variable sampling. "
+            "can be tuple of 3 values (start, end, decay)",
+        )
+        parser.add_argument(
+            "--quantizer-depth",
+            type=int,
+            help="number of quantizer layers",
+        )
+        parser.add_argument(
+            "--quantizer-factor",
+            type=int,
+            help="number of quantizer layers",
+        )
+        parser.add_argument(
+            "--get-code-distribution",
+            action='store_true',
+            help="whether to get the code distribution (for test)",
+        )
+
+        # relative pos enc
+        parser.add_argument(
+            "--relative-position-embedding",
+            action='store_true',
+            help="whether to use relative position embedding",
+        )
+        parser.add_argument(
+            "--num-buckets",
+            type=int,
+            default=320,
+            help="num of buckets for relative position embedding",
+        )
+        parser.add_argument(
+            "--max-distance",
+            type=int,
+            default=1280,
+            help="max distance for relative position embedding",
+        )
+        parser.add_argument(
+            "--encoder-max-relative-position",
+            type=int,
+            help="max distance for relative position embedding in encoder",
+        )
+        parser.add_argument(
+            "--decoder-max-relative-position",
+            type=int,
+            help="max distance for relative position embedding in decoder",
+        )
+
+        # hubert feature extractor
+        parser.add_argument(
+            "--conv-feature-layers",
+            type=str,
+            help= "string describing convolutional feature extraction "
+            "layers in form of a python list that contains "
+            "[(dim, kernel_size, stride), ...]",
+        )
+        parser.add_argument(
+            "--conv-bias",
+            action='store_true',
+            help="include bias in conv encoder",
+        )
+        parser.add_argument(
+            "--extractor-mode",
+            choices=["default", "layer_norm"],
+            help="mode for feature extractor. default has a single group "
+            "norm with d groups in the first conv block, whereas layer_norm "
+            "has layer norms in every block (meant to use with normalize=True)"
+        )
+
+        # others
+        parser.add_argument(
+            "--bert-init",
+            action='store_true',
+            help="initilize as bert",
+        )
+        parser.add_argument(
+            "--unb-enc-layer",
+            type=int,
+            default=-1,
+            help="which layer's output is used as the input of decoder",
+        )
+
+    # Encoder, Decoder
+    @classmethod
+    def build_encoder(cls, args, dictionary=None, embed_tokens=None):
+        return TransformerEncoder(args, dictionary, embed_tokens)
+
+    @classmethod
+    def build_decoder(cls, args):
+        return TransformerDecoder(args)
+
+    # Encoder Prenet
+    @classmethod
+    def build_text_encoder_prenet(cls, embed_tokens, args):
+        return TextEncoderPrenet(embed_tokens, args)
+
+    @classmethod
+    def build_speech_encoder_prenet(cls, args):
+        return SpeechEncoderPrenet(args)
+
+    # Decoder Prenet
+    @classmethod
+    def build_text_decoder_prenet(cls, embed_tokens, args):
+        return TextDecoderPrenet(embed_tokens, args)
+
+    @classmethod
+    def build_speech_decoder_prenet(cls, odim, args):
+        return SpeechDecoderPrenet(odim, args)
+
+    # Decoder Postnet
+    @classmethod
+    def build_text_decoder_postnet(cls, embed_tokens, dictionary, args):
+        return TextDecoderPostnet(embed_tokens, dictionary, args)
+
+    @classmethod
+    def build_speaker_decoder_postnet(cls, embed_dim, class_num, args):
+        return SpeakerDecoderPostnet(embed_dim, class_num, args)
+
+    @classmethod
+    def build_speech_decoder_postnet(cls, odim, args):
+        return SpeechDecoderPostnet(odim, args)
+
+    @classmethod
+    def build_speech_encoder_postnet(cls, dictionaries, args):
+        return SpeechEncoderPostnet(dictionaries, args)
+
+    @classmethod
+    def build_model(cls, args, task):
+        """Build a new model instance."""
+
+        # make sure all arguments are present in older models
+        base_architecture(args)
+
+        def build_embedding(dictionary, embed_dim, max_num_embeddings=None):
+            num_embeddings = len(dictionary)
+            if max_num_embeddings is not None and isinstance(max_num_embeddings, int):
+                num_embeddings = min(num_embeddings, max_num_embeddings)  
+            padding_idx = dictionary.pad()
+            return Embedding(num_embeddings, embed_dim, padding_idx)
+
+        if hasattr(args, "sid_pad_prenet") and args.sid_pad_prenet:
+            max_num_embeddings = 3 # <pad> at index 2
+        else:
+            max_num_embeddings = None
+        
+        text_decoder_embed_tokens = build_embedding(
+            task.dicts["text"], args.decoder_embed_dim, max_num_embeddings
+        )        
+
+        if args.share_input_output_embed:
+            text_encoder_embed_tokens = text_decoder_embed_tokens
+        else:
+            text_encoder_embed_tokens = build_embedding(
+                task.dicts["text"], args.encoder_embed_dim
+            )
+
+        speech_odim = args.speech_odim
+        if "text" in task.dicts:
+            encoder = cls.build_encoder(args, task.dicts["text"], text_encoder_embed_tokens)
+        else:
+            encoder = cls.build_encoder(args)      
+        decoder = cls.build_decoder(args)
+
+        text_encoder_prenet = cls.build_text_encoder_prenet(text_encoder_embed_tokens, args)
+        speech_encoder_prenet = cls.build_speech_encoder_prenet(args)
+
+        text_decoder_prenet = cls.build_text_decoder_prenet(text_decoder_embed_tokens, args)
+        if getattr(args, "sid_pooling_layer", None) == "decoder-las":
+            speech_decoder_prenet = cls.build_speech_encoder_prenet(args)
+        else:
+            speech_decoder_prenet = cls.build_speech_decoder_prenet(speech_odim, args)
+
+        text_decoder_postnet = cls.build_text_decoder_postnet(text_decoder_embed_tokens, task.dicts['text'], args)
+        speech_decoder_postnet = cls.build_speech_decoder_postnet(speech_odim, args)
+
+        if getattr(args, "sid_t5_postnet", False):
+            speaker_decoder_postnet = None
+        else:
+            if task.t5_task == "s2c":
+                speaker_decoder_postnet = cls.build_speaker_decoder_postnet(args.sid_embed_dim, len(task.dicts['text']), args)
+            else:
+                speaker_decoder_postnet = None
+
+        if "hubert" in task.dicts:
+            speech_encoder_postnet = cls.build_speech_encoder_postnet(task.dicts['hubert'], args)
+        else:
+            speech_encoder_postnet = None
+
+        return cls(
+            args, 
+            encoder, decoder, 
+            text_encoder_prenet, speech_encoder_prenet,
+            text_decoder_prenet, speech_decoder_prenet,
+            text_decoder_postnet, speech_decoder_postnet,
+            speaker_decoder_postnet, speech_encoder_postnet,
+        )
+
+    def get_normalized_probs(
+        self,
+        net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]],
+        log_probs: bool,
+        sample: Optional[Dict[str, Tensor]] = None,
+    ):
+        # net_output['encoder_out'] is a (B, T, D) tensor
+        lprobs = self.get_normalized_probs_scriptable(net_output, log_probs, sample)
+        lprobs.batch_first = True
+        return lprobs
+
+    def get_normalized_probs_for_ctc(self, net_output, log_probs):
+        """Get normalized probabilities (or log probs) from a net's output."""
+
+        logits = net_output["encoder_out_for_ctc"][0]
+        if log_probs:
+            return utils.log_softmax(logits.float(), dim=-1)
+        else:
+            return utils.softmax(logits.float(), dim=-1)
+
+    def get_logits(self, net_output, is_masked=True):
+        if is_masked:
+            logits_list = net_output["logit_m_list"]
+        else:
+            logits_list = net_output["logit_u_list"]
+        logits_list = [x.float() for x in logits_list if x is not None]
+        return logits_list
+
+    def get_targets(self, sample, net_output, is_masked=True):
+        if "logit_m_list" in net_output:
+            logits_list = self.get_logits(net_output, is_masked)
+            targets_list = [
+                x.new_zeros(x.size(0), dtype=torch.long) for x in logits_list
+            ]
+            return targets_list
+        else:
+            return sample["target"]
+
+    def get_extra_losses(self, net_output):
+        extra_losses = []
+        names = []
+
+        if "features_pen" in net_output:
+            extra_losses.append(net_output["features_pen"])
+            names.append("features_pen")
+
+        if "prob_perplexity" in net_output:
+            extra_losses.append(
+                (net_output["num_vars"] - net_output["prob_perplexity"])
+                / net_output["num_vars"]
+            )
+            names.append("prob_perplexity")
+
+        return extra_losses, names
+
+    def forward(self, source=None, src_tokens=None, src_lengths=None, prev_output_tokens=None, tgt_lengths=None, spkembs=None, target_list=None, task_name=None, padding_mask=None, only_hubert=False, only_ctc=False, feature_only=False, tgt_enc_layer=None, mask=True):
+        """
+        The forward method inherited from the base class has a **kwargs
+        argument in its input, which is not supported in torchscript. This
+        method overwrites the forward method definition without **kwargs.
+        """
+        assert source is not None or src_tokens is not None
+        # padding_mask is not none only when input is waveform
+        if source is None and padding_mask is None and not feature_only:
+            input_type = 'text'
+        else:
+            input_type = 'speech'
+
+        if prev_output_tokens is not None and len(prev_output_tokens.size()) == 2:
+            output_type = 'text'
+            codebook_out = {}
+        else:
+            output_type = 'speech'
+
+        if task_name is not None and task_name == "s2c":
+            if target_list is not None and target_list.size(1) == 1 and not getattr(self.args, "sid_t5_postnet", False):
+                sid_target = F.one_hot(target_list.squeeze(1), num_classes=self.speaker_decoder_postnet.class_num)
+            else:
+                sid_target = None
+            target_list = None
+
+        # Encoder Prenet
+        if input_type == 'text':
+            encoder_input, encoder_padding_mask = self.text_encoder_prenet(src_tokens)
+        else:
+            if target_list is not None:
+                encoder_input, encoder_padding_mask = self.speech_encoder_prenet(source, require_feat_pen=True, target_list=target_list, padding_mask=padding_mask, mask=mask)
+                encoder_input, features_pen, mask_indices, target_list = encoder_input
+            else:
+                encoder_input, encoder_padding_mask = self.speech_encoder_prenet(source, padding_mask=padding_mask, mask=self.training)
+                # shuffle a batch of inputs of encoder
+                if self.training and hasattr(self.args, "sid_shuffle_encoder_input") and getattr(self.args, "sid_shuffle_encoder_input", False):
+                    shuffle_index = torch.randperm(encoder_padding_mask.size(1), device=encoder_padding_mask.device)
+                    encoder_input = torch.index_select(encoder_input, 1, shuffle_index)
+                    encoder_padding_mask = torch.index_select(encoder_padding_mask, 1, shuffle_index)
+                if getattr(self.args, "sid_encoder_cls", None) == "encoder":
+                    prev_output_tokens = torch.zeros_like(prev_output_tokens)
+                    encoder_input, encoder_padding_mask = self._integrate_with_speaker_cls(prev_output_tokens, encoder_input, encoder_padding_mask)
+
+        # Encoder: T x B x C
+        encoder_output = self.encoder(encoder_input, encoder_padding_mask, tgt_layer=tgt_enc_layer)
+
+        if task_name is not None and task_name == 'speech_pretrain' and feature_only:
+            return encoder_output["encoder_out"][0].transpose(0, 1)
+
+        if task_name is not None and task_name == 's2c':
+            if self.args.sid_pooling_layer == "encoder":
+                return self.speaker_decoder_postnet(encoder_output["encoder_out"][0].transpose(0, 1).mean(1), sid_target), None
+            elif self.args.sid_pooling_layer == "encoder-cls":
+                return self.speaker_decoder_postnet(encoder_output["encoder_out"][0].transpose(0, 1)[:,0], sid_target), None
+            elif self.args.sid_pooling_layer == "encoder-speaker" or getattr(self.args, "sid_decoder_speaker", False):
+                return self.speaker_decoder_postnet(encoder_output["encoder_out"][0].transpose(0, 1), sid_target), None
+
+        if target_list is not None:
+            hubert_results = self.hubert_layer(
+                encoder_output["encoder_out"][0].transpose(0, 1), 
+                encoder_padding_mask, 
+                mask_indices, 
+                target_list
+            )
+
+            hubert_results['features_pen'] = features_pen
+
+        if "decoder_input" in encoder_output and encoder_output["decoder_input"][0] is not None:
+            # Change the encoder output to decoder input once set unb-enc-layer
+            encoder_output["encoder_out"] = encoder_output["decoder_input"]
+
+        if self.use_codebook:
+            q = self.quantizer(encoder_output["encoder_out"][0].transpose(0, 1))
+
+            # q["x"]: B x T x C
+            # Sample indexs according to the codebook prob
+            random_idx = torch.randperm(q["x"].size(1))[:int(q["x"].size(1) * self.codebook_prob)]
+            # Make weight for q
+            q_w = q["x"].new_zeros(q["x"].size(1))
+            q_w[random_idx] = 1.0
+            # Combine quantized codes and encoder output
+            encoder_output["encoder_out"][0] = (
+                q_w.view(-1, 1) * q["x"] + (- q_w + 1).view(-1, 1) * encoder_output["encoder_out"][0].transpose(0, 1)
+            ).transpose(0, 1)
+
+            # encoder_output["encoder_out"][0] = q["x"].transpose(0, 1)
+            if output_type == 'speech':
+                hubert_results["prob_perplexity"] = q["prob_perplexity"]
+                hubert_results["code_perplexity"] = q["code_perplexity"]
+                hubert_results["num_vars"] = q["num_vars"]
+                hubert_results["temp"] = q["temp"]
+            elif output_type == 'text':
+                codebook_out["prob_perplexity"] = q["prob_perplexity"]
+                codebook_out["code_perplexity"] = q["code_perplexity"]
+                codebook_out["num_vars"] = q["num_vars"]
+                codebook_out["temp"] = q["temp"]
+
+        if only_hubert and target_list is not None:
+            return hubert_results, None
+        
+        if only_ctc and task_name is not None and task_name == "s2t":
+            return None, encoder_output
+        elif not self.training and prev_output_tokens is None and task_name == "s2t" and task_name is not None:
+            return encoder_output
+
+        # Decoder Prenet
+        if output_type == 'text':
+            # _ is the incremental state
+            prev_output_tokens, tgt_mask, _ = self.text_decoder_prenet(prev_output_tokens)
+            if task_name is not None and task_name == 's2c':
+                prev_output_tokens = torch.zeros_like(prev_output_tokens)
+        else:
+            # integrate speaker embedding
+            if self.spk_embed_integration_type == "pre" and self.spk_embed_dim is not None:
+                # Decoder Prenet
+                prev_output_tokens, tgt_mask = self.speech_decoder_prenet(prev_output_tokens, tgt_lengths, spkembs)
+            else:
+                if self.spk_embed_dim is not None:
+                    encoder_output["encoder_out"] = [self._integrate_with_spk_embed(
+                        encoder_output["encoder_out"][0].transpose(0, 1), spkembs
+                    ).transpose(0, 1)]
+
+                prev_output_tokens, tgt_mask = self.speech_decoder_prenet(prev_output_tokens, tgt_lengths)
+
+        # BART Sequence Classification: cat <pad> + feature before decoder
+        if task_name is not None and task_name == 's2c' and self.args.sid_pooling_layer == "decoder-las":
+            decoder_feat_input, decoder_feat_mask = self.speech_decoder_prenet(src_tokens, src_lengths)
+            prev_output_tokens, tgt_mask = self._integrate_with_speaker_cls((prev_output_tokens, tgt_mask), decoder_feat_input, decoder_feat_mask, cls_first=False)
+        
+        # SE predict masking to corresponding inputs and source speech replaces the prev_output_tokens as the input of decoder
+        if task_name is not None and task_name == "s2s" and getattr(self.args, "se_decoder_input", "previous_target") == "source":
+            prev_output_tokens, tgt_mask = self.speech_decoder_prenet(src_tokens, src_lengths)
+
+        # Decoder
+        decoder_output, extra = self.decoder(prev_output_tokens, tgt_mask, encoder_output, 
+                                             full_context_alignment=getattr(self.args, "decoder_full_context_alignment", False), 
+                                             alignment_layer=(-1 if target_list is None and output_type == 'speech' else None))
+        # Decoder Postnet
+        if task_name is not None and task_name == 's2c':
+            if not getattr(self.args, "sid_t5_postnet", False):
+                if self.args.sid_pooling_layer == "decoder":
+                    return self.speaker_decoder_postnet(decoder_output.mean(1), sid_target), None
+                elif self.args.sid_pooling_layer == "decoder-las":
+                    indices = (tgt_mask.eq(False).float().sum(1) - 1.0).type(torch.int64)
+                    indices = indices.unsqueeze(1).unsqueeze(2).expand(-1, -1, decoder_output.size(2))
+                    return self.speaker_decoder_postnet(decoder_output.gather(1, indices), sid_target), None
+            else:
+                return (self.text_decoder_postnet(decoder_output), None), encoder_output
+
+        # SE predict: masking, target, delta. Ensure reduction factor 1
+        if task_name is not None and task_name == 's2s' and getattr(self.args, "se_predict", None) is not None:
+            assert self.reduction_factor == 1, f"{self.reduction_factor} != 1"
+            before_outs, after_outs, logits = self.speech_decoder_postnet(decoder_output)
+            se_predict = getattr(self.args, "se_predict")
+            if se_predict == "masking":
+                before_outs = torch.sigmoid(before_outs) * src_tokens
+                after_outs = torch.sigmoid(after_outs) * src_tokens
+                return before_outs, after_outs, logits, extra['attn'][0]
+            elif se_predict == "target":
+                return before_outs, after_outs, logits, extra['attn'][0]
+            elif se_predict == "delta":
+                before_outs = before_outs - src_tokens
+                after_outs = after_outs - src_tokens
+                return before_outs, after_outs, logits, extra['attn'][0]
+            else:
+                raise ValueError(f"{se_predict} not in [masking, target, delta]")
+
+        if task_name is not None and task_name == 's2t':
+            #return self.text_decoder_postnet(decoder_output), None
+            return (self.text_decoder_postnet(decoder_output), None), encoder_output
+        if output_type == 'text':
+            return (self.text_decoder_postnet(decoder_output), None), codebook_out, encoder_output
+        else:
+            if target_list is not None:
+                return hubert_results, (self.speech_decoder_postnet(decoder_output) + (extra['attn'][0],))
+            else:
+                return self.speech_decoder_postnet(decoder_output) + (extra['attn'][0],)
+
+    def _integrate_with_speaker_cls(self, pad_input, encoder_input, encoder_padding_mask=None, cls_first=True):
+        """
+        encoder_input: [B, T, C]
+        encoder_padding_mask: [B, T]
+        """
+        if hasattr(self, "text_decoder_prenet"):
+            if isinstance(pad_input, tuple):
+                repeat_cls_vector, repeat_cls_mask = pad_input
+            else:
+                repeat_cls_vector, repeat_cls_mask, _ = self.text_decoder_prenet(pad_input)
+
+            if encoder_padding_mask is not None:
+                bsz = encoder_input.size(0)
+                tsz = encoder_input.size(1)
+                encoder_padding_mask = encoder_input.new_zeros((bsz, tsz)) == 1.0
+            if repeat_cls_mask is None:
+                mask_size = (encoder_padding_mask.size(0), 1)
+                mask_type = encoder_padding_mask.dtype
+                repeat_cls_mask = encoder_padding_mask.new_zeros(mask_size) == 1.0
+            ret_encoder_padding_mask = torch.cat([repeat_cls_mask, encoder_padding_mask], dim=1)
+
+            if cls_first:
+                ret_encoder_input = torch.cat([repeat_cls_vector, encoder_input], dim=1)
+            else:
+                ret_encoder_input = torch.cat([encoder_input, encoder_input[:,-1:,:]], dim=1)
+                mask_size = (encoder_padding_mask.size(0), 1)
+                mask_type = encoder_padding_mask.dtype
+                repeat_cls_mask_ = encoder_padding_mask.new_ones(mask_size) == 1.0
+                encoder_padding_mask_ = torch.cat([encoder_padding_mask, repeat_cls_mask_], dim=1)
+                indices = encoder_padding_mask.eq(False).float().sum(1).type(torch.int64).unsqueeze(1)
+                indices_mask = torch.zeros_like(ret_encoder_padding_mask).scatter(1, indices, 1.0)
+                ret_encoder_input = ret_encoder_input * (1.0 - encoder_padding_mask_.type(ret_encoder_input.dtype).unsqueeze(2)) \
+                    + repeat_cls_vector * indices_mask.type(repeat_cls_vector.dtype).unsqueeze(2)
+            
+        return ret_encoder_input, ret_encoder_padding_mask
+
+    def _integrate_with_spk_embed(self, hs, spembs):
+        """Integrate speaker embedding with hidden states.
+        Args:
+            hs (Tensor): Batch of hidden state sequences (B, Tmax, adim).
+            spembs (Tensor): Batch of speaker embeddings (B, spk_embed_dim).
+        Returns:
+            Tensor: Batch of integrated hidden state sequences (B, Tmax, adim)
+        """
+        if self.spk_embed_integration_type == "add":
+            # apply projection and then add to hidden states
+            spembs = self.projection(F.normalize(spembs))
+            hs = hs + spembs.unsqueeze(1)
+        elif self.spk_embed_integration_type == "concat":
+            # concat hidden states with spk embeds and then apply projection
+            spembs = F.normalize(spembs).unsqueeze(1).expand(-1, hs.size(1), -1)
+            hs = self.projection(torch.cat([hs, spembs], dim=-1))
+        else:
+            raise NotImplementedError("support only add or concat.")
+
+        return hs
+
+    def load_state_dict(
+        self,
+        state_dict,
+        strict=True,
+        model_cfg=None,
+        args=None,
+    ):
+        """NOT STRICT Copies parameters and buffers from *state_dict* into this module and
+        its descendants.
+
+        Overrides the method in :class:`nn.Module`. Compared with that method
+        this additionally "upgrades" *state_dicts* from old checkpoints.
+        """
+        # self.prune_modules(model_cfg.modules_filter)
+        model_dict_size = self.text_decoder_postnet.output_projection.out_features
+        ckpt_dict_size = state_dict["text_decoder_postnet.output_projection.weight"].size(0)
+        if model_dict_size != ckpt_dict_size:
+            # reset dictionary-related modules, such as embedding table and encoder ctc embed
+            logger.warn(f"not equal dictionary between model and checkpoint: {model_dict_size} vs {ckpt_dict_size}")
+            logger.info(f"reset model dictionary with size of {model_dict_size}")
+            removed_keys = [
+                key for key in state_dict.keys() if any(
+                    key.startswith(previ) for previ in [
+                        "encoder.proj", "text_encoder_prenet", "text_decoder_prenet", "text_decoder_postnet"
+                    ]
+                )
+            ]
+            for key in removed_keys:
+                state_dict.pop(key, None)
+                logger.info(f"removed loaded checkpoint: {key}")
+        for m in self._modules.keys():
+            m_state_dict = {
+                key.replace(f"{m}.", ""): value for key, value in state_dict.items() if key.startswith(f"{m}.")
+            }
+            if hasattr(self, m):
+                self._modules[m].load_state_dict(m_state_dict, False)
+        return self
+
+    def prune_modules(self, modules_filter=None):
+        """Prune unused modules for specific tasks."""
+        if modules_filter is None:
+            return
+        elif modules_filter == "s2c":
+            if hasattr(self, "text_encoder_prenet"): del self.text_encoder_prenet
+            if hasattr(self, "speech_decoder_prenet") and getattr(self.args, "sid_pooling_layer", None) != "decoder-las": 
+                del self.speech_decoder_prenet
+            if hasattr(self, "speech_decoder_postnet"): del self.speech_decoder_postnet
+            if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
+            if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
+            if hasattr(self.encoder, "proj"): self.encoder.proj = None
+            if hasattr(self, "projection"): del self.projection
+            if hasattr(self, "quantizer"): del self.quantizer
+            if getattr(self.args, "sid_pooling_layer", "decoder").startswith("encoder") or getattr(self.args, "sid_decoder_speaker", False): 
+                if hasattr(self.decoder, "dropout_module"): del self.decoder.dropout_module
+                if hasattr(self.decoder, "layers"): del self.decoder.layers
+                if hasattr(self.decoder, "layer_norm"): del self.decoder.layer_norm
+                if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
+        elif modules_filter == "s2s":
+            if hasattr(self, "speaker_decoder_postnet"): del self.speaker_decoder_postnet
+            if hasattr(self, "text_encoder_prenet"): del self.text_encoder_prenet
+            if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
+            if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
+            if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
+            if hasattr(self.encoder, "proj"): self.encoder.proj = None
+            if hasattr(self, "projection"): del self.projection
+            if hasattr(self, "quantizer"): del self.quantizer
+        elif modules_filter == "t2s":
+            if hasattr(self, "speaker_decoder_postnet"): del self.speaker_decoder_postnet
+            if hasattr(self, "speech_encoder_prenet"): del self.speech_encoder_prenet
+            if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
+            if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
+            if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
+            if hasattr(self.encoder, "proj"): self.encoder.proj = None
+            if hasattr(self, "projection"): del self.projection
+            if hasattr(self, "quantizer"): del self.quantizer
+        elif modules_filter == "s3prl":
+            # remain the encoder and the pre/post net
+            if hasattr(self.decoder, "dropout_module"): del self.decoder.dropout_module
+            if hasattr(self.decoder, "layers"): del self.decoder.layers
+            if hasattr(self.decoder, "layer_norm"): del self.decoder.layer_norm
+            if hasattr(self, "speaker_decoder_postnet"): del self.speaker_decoder_postnet
+            if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
+            if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
+            if hasattr(self, "speech_decoder_prenet"): del self.speech_decoder_prenet
+            if hasattr(self, "speech_decoder_postnet"): del self.speech_decoder_postnet
+            if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
+            if hasattr(self.encoder, "proj"): self.encoder.proj = None
+            if hasattr(self, "projection"): del self.projection
+            if hasattr(self, "quantizer"): del self.quantizer
+
+    def forward_encoder_torchscript(self, net_input: Dict[str, Tensor]):
+        """A TorchScript-compatible version of forward.
+
+        Encoders which use additional arguments may want to override
+        this method for TorchScript compatibility.
+        """
+        if torch.jit.is_scripting():
+            return self.forward_encoder(
+                source=net_input["source"],
+                padding_mask=net_input["padding_mask"]
+            )
+        else:
+            return self.forward_encoder_non_torchscript(net_input)
+
+    @torch.jit.unused
+    def forward_encoder_non_torchscript(self, net_input: Dict[str, Tensor]):
+        encoder_input = {
+            k: v for k, v in net_input.items() if k != "prev_output_tokens" and k != "task_name"
+        }
+        return self.forward_encoder(**encoder_input)
+
+    def forward_encoder(self, source, padding_mask=None):
+        # Encoder Prenet
+        encoder_input, encoder_padding_mask = self.speech_encoder_prenet(source, padding_mask=padding_mask, mask=False)
+
+        # Encoder
+        encoder_output = self.encoder(encoder_input, encoder_padding_mask)
+
+        return encoder_output
+
+    def forward_text_encoder(self, src_tokens):
+        # Text Encoder Prenet
+        encoder_input, encoder_padding_mask = self.text_encoder_prenet(src_tokens)
+
+        # Encoder
+        encoder_output = self.encoder(encoder_input, encoder_padding_mask)
+
+        return encoder_output
+
+    def forward_decoder(self, tokens, encoder_out, incremental_state):
+        # Decoder Prenet
+        prev_output_tokens, tgt_mask, incremental_state = self.text_decoder_prenet(tokens, incremental_state)
+
+        # Decoder
+        decoder_output, extra = self.decoder(
+            prev_output_tokens,
+            tgt_mask,
+            encoder_out=encoder_out,
+            incremental_state=incremental_state,
+        )
+
+        # Decoder Postnet
+        return self.text_decoder_postnet(decoder_output), extra
+
+    def set_num_updates(self, num_updates):
+        """Set the number of parameters updates."""
+        super().set_num_updates(num_updates)
+        self.num_updates = num_updates
+
+    def generate_class(self, source, prev_output_tokens, **kwargs):
+        encoder_out = self.forward_encoder(source, padding_mask=kwargs["padding_mask"])
+
+        prev_output_tokens, tgt_mask, _ = self.text_decoder_prenet(prev_output_tokens, {})
+        prev_output_tokens = torch.zeros_like(prev_output_tokens) # s2c use zero vector as [CLS]
+
+        decoder_output, extra = self.decoder(
+            prev_output_tokens,
+            tgt_mask,
+            encoder_out=encoder_out,
+        )
+
+        decoder_out, embed = self.speaker_decoder_postnet(decoder_output.mean(1))
+
+        pred_class = decoder_out.argmax(1)
+        return pred_class
+
+    def generate_speech(self, source=None, src_tokens=None, spkembs=None, **kwargs):
+        assert source is not None or src_tokens is not None
+
+        threshold = kwargs.get("threshold", 0.5)
+        minlenratio = kwargs.get("threshold", 0.0)
+
+        if source is None:
+            assert src_tokens.size(0) == 1
+            encoder_out = self.forward_text_encoder(src_tokens)
+            maxlenratio = kwargs.get("threshold", 20.0)
+        else:
+            assert source.size(0) == 1
+            encoder_out = self.forward_encoder(source, padding_mask=kwargs["padding_mask"])
+            maxlenratio = kwargs.get("threshold", 10.0)
+
+        if spkembs is not None and self.spk_embed_integration_type != "pre":
+            encoder_out["encoder_out"] = [self._integrate_with_spk_embed(
+                encoder_out["encoder_out"][0].transpose(0, 1), spkembs
+            ).transpose(0, 1)]
+            spkembs = None
+
+        maxlen = int(encoder_out["encoder_out"][0].size(0) * maxlenratio / self.reduction_factor)
+        minlen = int(encoder_out["encoder_out"][0].size(0) * minlenratio / self.reduction_factor)
+        
+        idx = 0
+        ys = encoder_out["encoder_out"][0].new_zeros(1, 1, self.speech_decoder_postnet.odim)
+        outs, probs = [], []
+
+        # forward decoder step-by-step
+        if isinstance(self.decoder, FairseqIncrementalDecoder):
+            incremental_states = {}
+        else:
+            incremental_states = None
+        attns = []
+        while True:
+            # update index
+            idx += 1
+            # calculate output and stop prob at idx-th step
+            decoder_in, _ = self.speech_decoder_prenet(ys, spkembs=spkembs)
+            z, extra = self.decoder(decoder_in[:,-1:], None, encoder_out, incremental_states, alignment_layer=-1)
+            outs += [self.speech_decoder_postnet.feat_out(z[0, -1]).view(self.reduction_factor, self.speech_decoder_postnet.odim)]  # [(r, odim), ...]
+            probs += [torch.sigmoid(self.speech_decoder_postnet.prob_out(z[0, -1]))]  # [(r), ...]
+
+            # update next inputs
+            ys = torch.cat((ys, outs[-1][-1].view(1, 1, self.speech_decoder_postnet.odim)), dim=1)  # (1, idx + 1, odim)
+            attns.append(torch.stack([att_l[0] for att_l in extra['attn'][0]], dim=0))
+            # check whether to finish generation
+            if int(sum(probs[-1] >= threshold)) > 0 or idx >= maxlen:
+                # check mininum length
+                if idx < minlen:
+                    continue
+                outs = (torch.cat(outs, dim=0).unsqueeze(0).transpose(1, 2))  # (L, odim) -> (1, L, odim) -> (1, odim, L)
+                if self.speech_decoder_postnet.postnet is not None:
+                    outs = outs + self.speech_decoder_postnet.postnet(outs)  # (1, odim, L)
+                outs = outs.transpose(2, 1).squeeze(0)  # (L, odim)
+                probs = torch.cat(probs, dim=0)
+                attn = torch.cat(attns, dim=2)
+                break
+
+        if outs.size(0) == maxlen:
+            logging.warning("output length reaches maximum length")
+        return outs, probs, attn
+
+
+@register_model_architecture(model_name="artst_transformer", arch_name="artst_transformer")
+def base_architecture(args):
+    # Transformer
+    args.bert_init = getattr(args, "bert_init", False)
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 768)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 768 * 4)
+    args.encoder_layers = getattr(args, "encoder_layers", 12)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 12)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
+    args.decoder_ffn_embed_dim = getattr(
+        args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
+    )
+    args.decoder_layers = getattr(args, "decoder_layers", 6)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 12)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
+    args.dropout = getattr(args, "dropout", 0.1)
+    args.attention_dropout = getattr(args, "attention_dropout", args.dropout)
+    args.activation_dropout = getattr(args, "activation_dropout", args.dropout)
+    args.activation_fn = getattr(args, "activation_fn", "gelu")
+    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.0)
+    args.decoder_output_dim = getattr(
+        args, "decoder_output_dim", args.decoder_embed_dim
+    )
+    args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
+    args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0)
+    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0)
+    args.max_text_positions = getattr(args, "max_text_positions", DEFAULT_MAX_TEXT_POSITIONS)
+    args.max_speech_positions = getattr(args, "max_speech_positions", DEFAULT_MAX_SPEECH_POSITIONS)
+
+    # Espnet related, including prenet, postnet
+    args.eprenet_conv_layers = getattr(args, "eprenet_conv_layers", 0)
+    args.eprenet_conv_filts = getattr(args, "eprenet_conv_filts", 0)
+    args.eprenet_conv_chans = getattr(args, "eprenet_conv_chans", 0)
+    args.use_batch_norm = getattr(args, "use_batch_norm", True)
+    args.eprenet_dropout_rate = getattr(args, "eprenet_dropout_rate", 0.0)
+    args.enc_use_scaled_pos_enc = getattr(args, "enc_use_scaled_pos_enc", True)
+    args.dec_use_scaled_pos_enc = getattr(args, "dec_use_scaled_pos_enc", True)
+    args.postnet_layers = getattr(args, "postnet_layers", 5)
+    args.postnet_chans = getattr(args, "postnet_chans", 256)
+    args.postnet_filts = getattr(args, "postnet_filts", 5)
+    args.postnet_dropout_rate = getattr(args, "postnet_dropout_rate", 0.5)
+    args.dprenet_dropout_rate = getattr(args, "dprenet_dropout_rate", 0.5)
+    args.dprenet_layers = getattr(args, "dprenet_layers", 2)
+    args.dprenet_units = getattr(args, "dprenet_units", 256)
+    args.initial_encoder_alpha = getattr(args, "initial_encoder_alpha", 1.0)
+    args.initial_decoder_alpha = getattr(args, "initial_decoder_alpha", 1.0)
+    args.spk_embed_integration_type = getattr(args, "spk_embed_integration_type", "pre")
+    args.spk_embed_dim = getattr(args, "spk_embed_dim", 512)
+    args.encoder_reduction_factor = getattr(args, "encoder_reduction_factor", 1)
+    args.reduction_factor = getattr(args, "reduction_factor", 2)
+    args.transformer_enc_positional_dropout_rate = getattr(args, "transformer_enc_positional_dropout_rate", 0.1)
+    args.transformer_dec_positional_dropout_rate = getattr(args, "transformer_dec_positional_dropout_rate", 0.1)
+    args.layer_norm_eps = getattr(args, "layer_norm_eps", 1e-5)
+    args.no_scale_embedding = getattr(args, "no_scale_embedding", True)
+    # Convolutional subsampler
+    args.encoder_speech_prenet = getattr(args, "encoder_speech_prenet", "conv")
+    args.conv_kernel_sizes = getattr(args, "conv_kernel_sizes", "5,5")
+    args.conv_channels = getattr(args, "conv_channels", 1024)
+    args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)
+
+    args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
+    args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
+    args.no_token_positional_embeddings = getattr(
+        args, "no_token_positional_embeddings", False
+    )
+    args.adaptive_input = getattr(args, "adaptive_input", False)
+    args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
+    args.share_input_output_embed = getattr(args, "share_input_output_embed", False)
+    args.share_ctc_embed = getattr(args, "share_ctc_embed", False)
+    args.freeze_encoder_updates = getattr(args, "freeze_encoder_updates", 0)
+    args.freeze_decoder_updates = getattr(args, "freeze_decoder_updates", 0)
+    args.no_freeze_encoder_layer = getattr(args, "no_freeze_encoder_layer", None)
+
+    ## sid
+    args.sid_embed_dim = getattr(args, "sid_embed_dim", 128)
+    args.sid_pooling_layer = getattr(args, "sid_pooling_layer", "decoder")
+    args.softmax_scale = getattr(args, "softmax_scale", 1)
+    args.softmax_margin = getattr(args, "softmax_margin", 0)
+    args.softmax_easy_margin = getattr(args, "softmax_easy_margin", False)
+    args.modules_filter = getattr(args, "modules_filter", None)
+
+    ## Hubert
+    args.conv_pos = getattr(args, "conv_pos", 128)
+    args.conv_pos_groups = getattr(args, "conv_pos_groups", 16)
+    args.target_glu = getattr(args, "target_glu", False)
+    args.logit_temp = getattr(args, "logit_temp", 0.1)
+    args.final_dim = getattr(args, "final_dim", 256)
+    args.untie_final_proj = getattr(args, "untie_final_proj", True)
+    args.feature_grad_mult = getattr(args, "feature_grad_mult", 0.1)
+    args.use_sent_enc_layer = getattr(args, "use_sent_enc_layer", True)
+    # hubert feature extractor
+    args.extractor_mode = getattr(args, "extractor_mode", "default")
+    args.conv_feature_layers = getattr(args, "conv_feature_layers", "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2")
+    args.conv_bias = getattr(args, "conv_bias", False)
+    # mask
+    args.hubert_mask_length = getattr(args, "hubert_mask_length", 10)
+    args.mask_prob = getattr(args, "mask_prob", 0.0)
+    args.mask_selection = getattr(args, "mask_selection", "static")
+    args.mask_other = getattr(args, "mask_other", 0)
+    args.no_mask_overlap = getattr(args, "no_mask_overlap", False)
+    args.mask_min_space = getattr(args, "mask_min_space", 1)
+    # channel mask
+    args.mask_channel_length = getattr(args, "mask_channel_length", 10)
+    args.mask_channel_prob = getattr(args, "mask_channel_prob", 0.0)
+    args.mask_channel_selection = getattr(args, "mask_channel_selection", "static")
+    args.mask_channel_other = getattr(args, "mask_channel_other", 0)
+    args.no_mask_channel_overlap = getattr(args, "no_mask_channel_overlap", False)
+    args.mask_channel_min_space = getattr(args, "mask_channel_min_space", 1)
+    # loss computation
+    args.skip_masked = getattr(args, "skip_masked", False)
+    args.skip_nomask = getattr(args, "skip_nomask", False)
+    # conv Pos
+    args.use_conv_pos = getattr(args, "use_conv_pos", False)
+    args.use_sinc_pos = getattr(args, "use_sinc_pos", False)
+
+    # codebook
+    args.use_codebook = getattr(args, "use_codebook", False)
+    args.latent_vars = getattr(args, "latent_vars", 100)
+    args.latent_groups = getattr(args, "latent_groups", 2)
+    args.latent_dim = getattr(args, "latent_dim", 0)
+    args.latent_temp = getattr(args, "latent_temp", (2, 0.5, 0.999995))
+    args.quantizer_depth = getattr(args, "quantizer_depth", 1)
+    args.quantizer_factor = getattr(args, "quantizer_factor", 3)
+    args.codebook_prob = getattr(args, "codebook_prob", 0.5)
+
+    # Relative pos embed
+    args.relative_position_embedding = getattr(args, "relative_position_embedding", False)
+    args.num_buckets = getattr(args, "num_buckets", 320)
+    args.max_distance = getattr(args, "max_distance", 1280)
+    args.encoder_max_relative_position = getattr(args, "encoder_max_relative_position", 160)
+    args.decoder_max_relative_position = getattr(args, "decoder_max_relative_position", 160)
+
+@register_model_architecture("artst_transformer", "artst_transformer_base")
+def artst_transformer_base(args):
+    args.use_conv_pos = getattr(args, "use_conv_pos", True)
+    args.use_sinc_pos = getattr(args, "use_sinc_pos", True)
+    args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
+    args.layer_norm_first = getattr(args, "layer_norm_first", False)
+    args.relative_position_embedding = getattr(args, "relative_position_embedding", True)
+    args.dropout = getattr(args, "dropout", 0.1)
+    args.activation_dropout = getattr(args, "activation_dropout", 0.0)
+    args.attention_dropout = getattr(args, "attention_dropout", 0.1)
+    args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.05)
+    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.05)
+    args.mask_prob = getattr(args, "mask_prob", 0.80)
+    base_architecture(args)
+
+@register_model_architecture("artst_transformer", "artst_transformer_large")
+def artst_transformer_large(args):
+    args.use_conv_pos = getattr(args, "use_conv_pos", True)
+    args.use_sinc_pos = getattr(args, "use_sinc_pos", True)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", True)
+    args.layer_norm_first = getattr(args, "layer_norm_first", True)
+    args.relative_position_embedding = getattr(args, "relative_position_embedding", True)
+    args.dropout = getattr(args, "dropout", 0.0)
+    args.activation_dropout = getattr(args, "activation_dropout", 0.0)
+    args.attention_dropout = getattr(args, "attention_dropout", 0.0)
+    args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.0)
+    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.0)
+    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
+    args.encoder_layers = getattr(args, "encoder_layers", 24)
+    args.decoder_layers = getattr(args, "decoder_layers", 6)
+    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4096)
+    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
+    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
+    args.feature_grad_mult = getattr(args, "feature_grad_mult", 1.0)
+    args.extractor_mode = getattr(args, "extractor_mode", "layer_norm")
+    args.final_dim = getattr(args, "final_dim", 768)
+    args.mask_prob = getattr(args, "mask_prob", 0.80)
+    base_architecture(args)
+
+@register_model_architecture("artst_transformer", "artst_transformer_base_asr")
+def artst_transformer_base_asr(args):
+    args.use_conv_pos = getattr(args, "use_conv_pos", True)
+    args.use_sinc_pos = getattr(args, "use_sinc_pos", True)
+    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
+    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
+    args.layer_norm_first = getattr(args, "layer_norm_first", False)
+    args.relative_position_embedding = getattr(args, "relative_position_embedding", True)
+    args.dropout = getattr(args, "dropout", 0.1)
+    args.activation_dropout = getattr(args, "activation_dropout", 0.1)
+    args.attention_dropout = getattr(args, "attention_dropout", 0.1)
+    args.feature_grad_mult = getattr(args, "feature_grad_mult", 0.0)
+    args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.1)
+    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.1)
+    args.mask_prob = getattr(args, "mask_prob", 0.75)
+    args.mask_selection = getattr(args, "mask_selection", "static")
+    args.mask_channel_length = getattr(args, "mask_channel_length", 64)
+    args.mask_channel_prob = getattr(args, "mask_channel_prob", 0.5)
+    args.mask_channel_selection = getattr(args, "mask_channel_selection", "static")
+    args.max_text_positions = getattr(args, "max_text_positions", 600)
+    base_architecture(args)