seamless_communication/inference/translator.py

# 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,
                ),
            )