Upload seamless_communication/inference/translator.py with huggingface_hub

seamless_communication/inference/translator.py

@@ -0,0 +1,428 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# This source code is licensed under the license found in the
+# MIT_LICENSE file in the root directory of this source tree.
+
+import logging
+from dataclasses import dataclass
+from enum import Enum, auto
+from pathlib import Path
+from typing import List, Optional, Tuple, Union, cast
+
+import torch
+import torch.nn as nn
+from fairseq2.assets import asset_store
+from fairseq2.assets.card import AssetCard
+from fairseq2.data import Collater, SequenceData, StringLike
+from fairseq2.data.audio import AudioDecoder, WaveformToFbankConverter
+from fairseq2.data.text import TextTokenizer
+from fairseq2.memory import MemoryBlock
+from fairseq2.nn.padding import PaddingMask, get_seqs_and_padding_mask
+from fairseq2.typing import DataType, Device
+from torch import Tensor
+
+from seamless_communication.inference.generator import (
+    SequenceGeneratorOptions,
+    UnitYGenerator,
+)
+from seamless_communication.models.unity import (
+    UnitTokenizer,
+    UnitYModel,
+    UnitYNART2UModel,
+    UnitYT2UModel,
+    load_unity_model,
+    load_unity_text_tokenizer,
+    load_unity_unit_tokenizer,
+    unity_archs,
+)
+from seamless_communication.models.vocoder import load_vocoder_model
+from seamless_communication.toxicity import (
+    ETOXBadWordChecker,
+    load_etox_bad_word_checker,
+)
+from seamless_communication.toxicity.mintox import mintox_pipeline
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s -- %(name)s: %(message)s",
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Task(Enum):
+    S2ST = auto()
+    S2TT = auto()
+    T2ST = auto()
+    T2TT = auto()
+    ASR = auto()
+
+
+class Modality(Enum):
+    SPEECH = "speech"
+    TEXT = "text"
+
+
+@dataclass
+class BatchedSpeechOutput:
+    units: List[List[int]]
+    """The batched list of generated units."""
+
+    audio_wavs: List[Tensor]
+    """The batched list of audio waveforms."""
+
+    sample_rate: int = 16000
+    """Sample rate of the audio waveforms."""
+
+
+class Translator(nn.Module):
+    def __init__(
+        self,
+        model_name_or_card: Union[str, AssetCard],
+        vocoder_name_or_card: Union[str, AssetCard, None],
+        device: Device,
+        text_tokenizer: Optional[TextTokenizer] = None,
+        apply_mintox: bool = False,
+        dtype: DataType = torch.float16,
+        input_modality: Optional[Modality] = None,
+        output_modality: Optional[Modality] = None,
+    ):
+        super().__init__()
+
+        if isinstance(model_name_or_card, str):
+            model_name_or_card = asset_store.retrieve_card(model_name_or_card)
+
+        assert isinstance(model_name_or_card, AssetCard)
+
+        if input_modality or output_modality:
+            unity_config = unity_archs.get_config(
+                model_name_or_card.field("model_arch").as_(str)
+            )
+            # Skip loading the text encoder.
+            if input_modality == Modality.SPEECH:
+                unity_config.use_text_encoder = False
+            # Skip loading the T2U model.
+            if output_modality == Modality.TEXT:
+                unity_config.t2u_config = None
+            model_name_or_card.field("model_config").set(unity_config)
+
+        # Load the model.
+        if device == torch.device("cpu"):
+            dtype = torch.float32
+
+        self.model = load_unity_model(model_name_or_card, device=device, dtype=dtype)
+        self.model.eval()
+        assert isinstance(self.model, UnitYModel)
+
+        if text_tokenizer is None:
+            self.text_tokenizer: TextTokenizer = load_unity_text_tokenizer(
+                model_name_or_card
+            )
+        else:
+            self.text_tokenizer = text_tokenizer
+
+        self.unit_tokenizer: Optional[UnitTokenizer] = None
+        if self.model.t2u_model is not None:
+            self.unit_tokenizer = load_unity_unit_tokenizer(model_name_or_card)
+
+        self.bad_word_checker: Optional[ETOXBadWordChecker] = None
+        if apply_mintox:
+            self.bad_word_checker = load_etox_bad_word_checker("mintox")
+
+        self.apply_mintox = apply_mintox
+
+        self.device = device
+        self.decode_audio = AudioDecoder(dtype=torch.float32, device=device)
+        self.convert_to_fbank = WaveformToFbankConverter(
+            num_mel_bins=80,
+            waveform_scale=2**15,
+            channel_last=True,
+            standardize=True,
+            device=device,
+            dtype=dtype,
+        )
+        self.collate = Collater(
+            pad_value=self.text_tokenizer.vocab_info.pad_idx or 0, pad_to_multiple=2
+        )
+        self.vocoder = None
+        if vocoder_name_or_card is not None and (
+            output_modality is None or output_modality == Modality.SPEECH
+        ):
+            self.vocoder = load_vocoder_model(
+                vocoder_name_or_card, device=device, dtype=dtype
+            )
+            self.vocoder.eval()
+
+    @classmethod
+    def get_prediction(
+        cls,
+        model: UnitYModel,
+        text_tokenizer: TextTokenizer,
+        unit_tokenizer: Optional[UnitTokenizer],
+        seqs: Tensor,
+        padding_mask: Optional[PaddingMask],
+        input_modality: Modality,
+        output_modality: Modality,
+        tgt_lang: str,
+        text_generation_opts: SequenceGeneratorOptions,
+        unit_generation_opts: Optional[SequenceGeneratorOptions],
+        unit_generation_ngram_filtering: bool = False,
+        duration_factor: float = 1.0,
+        prosody_encoder_input: Optional[SequenceData] = None,
+    ) -> Tuple[List[StringLike], Optional[Tensor]]:
+        # We disregard unit generations opts for the NAR T2U decoder.
+        if output_modality != Modality.SPEECH or isinstance(
+            model.t2u_model, UnitYNART2UModel
+        ):
+            unit_generation_opts = None
+
+        generator = UnitYGenerator(
+            model,
+            text_tokenizer,
+            tgt_lang,
+            unit_tokenizer if output_modality == Modality.SPEECH else None,
+            text_opts=text_generation_opts,
+            unit_opts=unit_generation_opts,
+        )
+
+        return generator(
+            seqs,
+            padding_mask,
+            input_modality.value,
+            output_modality.value,
+            ngram_filtering=unit_generation_ngram_filtering,
+            duration_factor=duration_factor,
+            prosody_encoder_input=prosody_encoder_input,
+        )
+
+    @staticmethod
+    def get_modalities_from_task_str(task_str: str) -> Tuple[Modality, Modality]:
+        try:
+            task = Task[task_str.upper()]
+        except KeyError:
+            raise ValueError(f"Unsupported task: {task_str}")
+
+        if task == Task.S2ST:
+            return Modality.SPEECH, Modality.SPEECH
+        # ASR is treated as S2TT with src_lang == tgt_lang
+        elif task == Task.S2TT or task == Task.ASR:
+            return Modality.SPEECH, Modality.TEXT
+        elif task == Task.T2TT:
+            return Modality.TEXT, Modality.TEXT
+        else:
+            return Modality.TEXT, Modality.SPEECH
+
+    @torch.inference_mode()
+    def predict(
+        self,
+        input: Union[str, Tensor, SequenceData],
+        task_str: str,
+        tgt_lang: str,
+        src_lang: Optional[str] = None,
+        text_generation_opts: Optional[SequenceGeneratorOptions] = None,
+        unit_generation_opts: Optional[SequenceGeneratorOptions] = None,
+        spkr: Optional[int] = -1,
+        sample_rate: int = 16000,
+        unit_generation_ngram_filtering: bool = False,
+        duration_factor: float = 1.0,
+        prosody_encoder_input: Optional[SequenceData] = None,
+        src_text: Optional[StringLike] = None,
+    ) -> Tuple[List[StringLike], Optional[BatchedSpeechOutput]]:
+        """
+        The main method used to perform inference on all tasks.
+
+        :param input:
+            Either text or path to audio or audio Tensor.
+        :param task_str:
+            String representing the task.
+            Valid choices are "S2ST", "S2TT", "T2ST", "T2TT", "ASR"
+        :param tgt_lang:
+            Target language to decode into.
+        :param src_lang:
+            Source language of input, only required for T2ST, T2TT tasks.
+        :param text_generation_opts:
+            Text generation hyperparameters for incremental decoding.
+        :param unit_generation_opts:
+            Unit generation hyperparameters for incremental decoding.
+        :param spkr:
+            Speaker id for vocoder.
+        :param unit_generation_ngram_filtering:
+            If True, removes consecutive repeated ngrams
+            from the decoded unit output.
+        :param src_text:
+            Optional source transcript (obtained by ASR for instance). This is used for
+            applying mintox toxicity mitigation. If this is not specify and apply_mintox=True
+            then src_lang must be specified and ASR will be run on the audio source.
+
+        :returns:
+            - Batched list of Translated text.
+            - Translated BatchedSpeechOutput.
+        """
+        input_modality, output_modality = self.get_modalities_from_task_str(task_str)
+
+        if self.apply_mintox and not (src_lang is not None or src_text is not None):
+            raise ValueError(
+                "`src_lang` must be specified when `apply_mintox` is `True` or you need to specify src_text."
+            )
+
+        if isinstance(input, dict):
+            src = cast(SequenceData, input)
+        elif input_modality == Modality.SPEECH:
+            audio = input
+            if isinstance(audio, str):
+                with Path(audio).open("rb") as fb:
+                    block = MemoryBlock(fb.read())
+                decoded_audio = self.decode_audio(block)
+            else:
+                assert (
+                    audio.dim() <= 2
+                ), "The audio tensor can't be more than 2 dimensions."
+                if audio.dim() == 1:
+                    audio = audio.unsqueeze(1)
+                elif audio.dim() == 2 and audio.size(0) < audio.size(1):
+                    logger.warning(
+                        "Transposing audio tensor from (bsz, seq_len) -> (seq_len, bsz)."
+                    )
+                    audio = audio.transpose(0, 1)
+
+                decoded_audio = {
+                    "waveform": audio,
+                    "sample_rate": sample_rate,
+                    "format": -1,
+                }
+            src = self.collate(self.convert_to_fbank(decoded_audio))["fbank"]
+        else:
+            if src_lang is None:
+                raise ValueError("src_lang must be specified for T2ST, T2TT tasks.")
+
+            text = input
+            assert isinstance(text, str)
+
+            self.token_encoder = self.text_tokenizer.create_encoder(
+                task="translation", lang=src_lang, mode="source", device=self.device
+            )
+            src = self.collate(self.token_encoder(text))
+
+        assert isinstance(self.model, UnitYModel)
+
+        seqs, padding_mask = get_seqs_and_padding_mask(src)
+
+        if text_generation_opts is None:
+            text_generation_opts = SequenceGeneratorOptions(
+                beam_size=5, soft_max_seq_len=(1, 200)
+            )
+        if unit_generation_opts is None:
+            unit_generation_opts = SequenceGeneratorOptions(
+                beam_size=5, soft_max_seq_len=(25, 50)
+            )
+
+        texts, units = self.get_prediction(
+            self.model,
+            self.text_tokenizer,
+            self.unit_tokenizer,
+            seqs,
+            padding_mask,
+            input_modality,
+            output_modality,
+            tgt_lang,
+            text_generation_opts,
+            unit_generation_opts,
+            unit_generation_ngram_filtering=unit_generation_ngram_filtering,
+            duration_factor=duration_factor,
+            prosody_encoder_input=prosody_encoder_input,
+        )
+
+        if self.apply_mintox and task_str != Task.ASR.name:
+            if input_modality == Modality.SPEECH:
+                if src_text is not None:
+                    src_texts = [src_text]
+                else:
+                    src_texts, _, = self.predict(
+                        input=input,
+                        task_str=Task.ASR.name,
+                        tgt_lang=tgt_lang,
+                        src_lang=src_lang,
+                        text_generation_opts=text_generation_opts,
+                        unit_generation_opts=unit_generation_opts,
+                        spkr=spkr,
+                        sample_rate=sample_rate,
+                        unit_generation_ngram_filtering=unit_generation_ngram_filtering,
+                    )
+            else:
+                assert isinstance(input, str)
+
+                src_texts = [input]
+
+            assert src_lang is not None
+            assert self.unit_tokenizer is not None
+            assert self.bad_word_checker is not None
+
+            texts, units = mintox_pipeline(
+                model=self.model,
+                text_tokenizer=self.text_tokenizer,
+                unit_tokenizer=self.unit_tokenizer,
+                device=self.device,
+                src_lang=src_lang,
+                tgt_lang=tgt_lang,
+                model_input=src,
+                input_modality=input_modality,
+                output_modality=output_modality,
+                src_texts=src_texts,
+                original_texts=texts,
+                original_units=units,
+                unit_generation_ngram_filtering=unit_generation_ngram_filtering,
+                text_generation_opts=text_generation_opts,
+                unit_generation_opts=unit_generation_opts,
+                bad_word_checker=self.bad_word_checker,
+                duration_factor=duration_factor,
+                prosody_encoder_input=prosody_encoder_input,
+            )
+
+        if output_modality == Modality.TEXT:
+            return texts, None
+        else:
+            assert units is not None
+
+            if isinstance(self.model.t2u_model, UnitYT2UModel):
+                # Remove the lang token for AR UnitY since the vocoder doesn't need it
+                # in the unit sequence. tgt_lang is fed as an argument to the vocoder.
+                units = units[:, 1:]
+                duration_prediction = True
+            else:
+                # Vocoder duration predictions not required since the NAR
+                # T2U model already predicts duration in the units.
+                duration_prediction = False
+
+            audio_wavs = []
+            speech_units = []
+            for i in range(len(units)):
+                assert self.model.t2u_model is not None
+                unit_padding_mask = (
+                    units[i] != self.model.t2u_model.target_vocab_info.pad_idx
+                )
+                u = units[i][unit_padding_mask]
+                speech_units.append(u.tolist())
+
+            if self.vocoder is not None:
+                translated_audio_wav = self.vocoder(
+                    units, tgt_lang, spkr, dur_prediction=duration_prediction
+                )
+                for i in range(len(units)):
+                    padding_removed_audio_wav = translated_audio_wav[
+                        i,
+                        :,
+                        : int(
+                            translated_audio_wav.size(-1)
+                            * len(speech_units[i])
+                            / len(units[i])
+                        ),
+                    ].unsqueeze(0)
+                    audio_wavs.append(padding_removed_audio_wav)
+            return (
+                texts,
+                BatchedSpeechOutput(
+                    units=speech_units,
+                    audio_wavs=audio_wavs,
+                    sample_rate=sample_rate,
+                ),
+            )