File size: 13,657 Bytes
1d81c94 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 |
# 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.
from dataclasses import dataclass
from typing import List, Optional, Tuple
import torch
from fairseq2.data import SequenceData, StringLike
from fairseq2.data.text import TextTokenizer
from fairseq2.generation import (
BeamSearchSeq2SeqGenerator,
Seq2SeqGenerator,
SequenceToTextConverter,
StepProcessor,
)
from fairseq2.nn.padding import (
PaddingMask,
apply_padding_mask,
get_seqs_and_padding_mask,
pad_seqs,
)
from fairseq2.nn.utils.module import infer_device
from torch import Tensor
from seamless_communication.models.unity.model import (
UnitYModel,
UnitYT2UModel,
UnitYX2TModel,
)
from seamless_communication.models.unity.unit_tokenizer import (
UnitTokenDecoder,
UnitTokenizer,
)
def remove_consecutive_repeated_ngrams(
sequence: List[int], min_size: int = 1, max_size: int = 40
) -> List[int]:
assert 1 <= min_size <= max_size
drop_idx = set() # indices that will be dropped from the sequence
# start from the beginning, check if an ngram of size k (for k=max..min) is
# followed by its copy, if so delete the first one, and start over after
# the deleted ngram.
start = 0
while start < len(sequence):
for k in range(max_size, min_size - 1, -1):
if sequence[start : start + k] == sequence[start + k : start + k + k]:
drop_idx |= set(range(start, start + k))
start += k - 1 # assumes repeating subsequences don't overlap
break
start += 1
return [token for idx, token in enumerate(sequence) if idx not in drop_idx]
@dataclass
class SequenceGeneratorOptions:
"""Holds the options to pass to a sequence generator."""
beam_size: int = 5
"""The beam size."""
soft_max_seq_len: Tuple[int, int] = (1, 200)
"""The terms ``a`` and ``b`` of ``ax + b`` where ``x`` is the source
sequence length. The generated sequences (including prefix sequence) will
have the maximum length of ``min(hard_max_seq_len, ax + b)``. See also
``hard_max_seq_len``."""
hard_max_seq_len: int = 1024
"""The hard limit on maximum length of generated sequences."""
step_processor: Optional[StepProcessor] = None
"""The processor called at each generation step."""
unk_penalty: float = 0.0
"""The UNK symbol penalty, where values less than 0 produce more UNKs;
values greater than 0 produce fewer UNKs."""
len_penalty: float = 1.0
"""The length penalty, where values less than 1.0 favor shorter
sequences; values greater than 1.0 favor longer sequences."""
class UnitYGenerator:
"""Generates text translations and speech units from a UnitY model."""
model: UnitYModel
s2t_converter: SequenceToTextConverter
t2t_converter: Optional[SequenceToTextConverter]
unit_decoder: Optional[UnitTokenDecoder]
unit_prefix_indices: Optional[Tensor]
unit_generator: Optional[Seq2SeqGenerator]
def __init__(
self,
model: UnitYModel,
text_tokenizer: TextTokenizer,
target_lang: str,
unit_tokenizer: Optional[UnitTokenizer] = None,
text_opts: Optional[SequenceGeneratorOptions] = None,
unit_opts: Optional[SequenceGeneratorOptions] = None,
) -> None:
"""
:param model:
The UnitY model to use for generation.
:param text_tokenizer:
The text tokenizer to use.
:param unit_tokenizer:
The unit tokenizer to use.
:param target_lang:
The target language.
:param text_generator_opts:
The options to pass to the underlying text :class:`Seq2SeqGenerator`.
:param unit_generator_opts:
The options to pass to the underlying unit :class:`Seq2SeqGenerator`.
"""
model.eval()
self.model = model
if text_opts is None:
text_opts = SequenceGeneratorOptions()
if model.text_decoder is None:
raise ValueError(
"`UnitYGenerator` requires a text decoder, but the current UnitY model does not have one."
)
assert model.text_decoder_frontend is not None
assert model.final_proj is not None
s2t_model = UnitYX2TModel(
encoder_frontend=model.speech_encoder_frontend,
encoder=model.speech_encoder,
decoder_frontend=model.text_decoder_frontend,
decoder=model.text_decoder,
final_proj=model.final_proj,
target_vocab_info=model.target_vocab_info,
)
step_processors = []
if text_opts.step_processor is not None:
step_processors.append(text_opts.step_processor)
generator = BeamSearchSeq2SeqGenerator(
s2t_model,
beam_size=text_opts.beam_size,
max_gen_len=text_opts.soft_max_seq_len,
max_seq_len=text_opts.hard_max_seq_len,
echo_prompt=True,
step_processors=step_processors,
unk_penalty=text_opts.unk_penalty,
len_penalty=text_opts.len_penalty,
)
self.s2t_converter = SequenceToTextConverter(
generator, text_tokenizer, "translation", target_lang
)
if model.text_encoder is None:
self.t2t_generator = None
else:
assert model.text_encoder_frontend is not None
assert model.text_encoder is not None
t2t_model = UnitYX2TModel(
encoder_frontend=model.text_encoder_frontend,
encoder=model.text_encoder,
decoder_frontend=model.text_decoder_frontend,
decoder=model.text_decoder,
final_proj=model.final_proj,
target_vocab_info=model.target_vocab_info,
)
generator = BeamSearchSeq2SeqGenerator(
t2t_model,
beam_size=text_opts.beam_size,
max_gen_len=text_opts.soft_max_seq_len,
max_seq_len=text_opts.hard_max_seq_len,
echo_prompt=True,
step_processors=step_processors,
unk_penalty=text_opts.unk_penalty,
len_penalty=text_opts.len_penalty,
)
self.t2t_converter = SequenceToTextConverter(
generator, text_tokenizer, "translation", target_lang
)
self.unit_generator = None
self.unit_decoder = None
# Set up unit generator.
if unit_tokenizer is not None:
if model.t2u_model is None:
raise ValueError(
"`model` does not have a T2U sub-model when `unit_tokenizer` is not None."
)
self.unit_decoder = unit_tokenizer.create_decoder()
unit_encoder = unit_tokenizer.create_encoder(
lang=target_lang, device=infer_device(model.t2u_model)
)
self.unit_prefix_indices = unit_encoder.prefix_indices
if isinstance(self.model.t2u_model, UnitYT2UModel):
if unit_opts is None:
# Speech sequences are typically much longer than text sequences.
unit_opts = SequenceGeneratorOptions(
soft_max_seq_len=(25, 50), hard_max_seq_len=5000
)
step_processors = []
if unit_opts.step_processor is not None:
step_processors.append(unit_opts.step_processor)
self.unit_generator = BeamSearchSeq2SeqGenerator(
self.model.t2u_model,
beam_size=unit_opts.beam_size,
max_gen_len=unit_opts.soft_max_seq_len,
max_seq_len=unit_opts.hard_max_seq_len,
echo_prompt=True,
step_processors=step_processors,
unk_penalty=unit_opts.unk_penalty,
len_penalty=unit_opts.len_penalty,
)
@torch.inference_mode()
def __call__(
self,
source_seqs: Tensor,
source_padding_mask: Optional[PaddingMask],
input_modality: str = "speech",
output_modality: str = "speech",
ngram_filtering: bool = False,
duration_factor: float = 1.0,
prosody_encoder_input: Optional[SequenceData] = None,
) -> Tuple[List[StringLike], Optional[Tensor]]:
"""
:param source_seqs:
The source sequences to use for generation. *Shape:* :math:`(N,S,*)`,
where :math:`N` is the batch size, :math:`S` is the sequence length,
and :math:`*` is any number of sequence-specific dimensions
including none.
:param source_padding_mask:
The padding mask of ``source_seqs``. *Shape:* :math:`(N,S)`, where
:math:`N` is the batch size and :math:`S` is the sequence length.
:param input_modality:
The type of modality to encode.
:param output_modality:
The type of modality to decode.
:param ngram_filtering:
If True, removes consecutive repeated ngrams
from the decoded unit output.
:returns:
- The output of the text generator.
- The output of the unit generator.
"""
if input_modality == "speech":
texts, text_gen_output = self.s2t_converter.batch_convert(
source_seqs, source_padding_mask
)
elif input_modality == "text":
if self.t2t_converter is None:
raise ValueError(
"Please set `use_text_encoder` to `True` in your model config to encode text."
)
texts, text_gen_output = self.t2t_converter.batch_convert(
source_seqs, source_padding_mask
)
else:
raise ValueError(f"Unsupported input_modality: {input_modality}")
# We skip T2U when we only need to output text.
if output_modality == "text":
return texts, None
assert self.model.target_vocab_info.pad_idx is not None
text_seq_list = [h[0].seq for h in text_gen_output.hypotheses]
text_seqs, text_padding_mask = pad_seqs(
text_seq_list, self.model.target_vocab_info.pad_idx
)
# Manually trim the final EOS token to be consistent with fairseq.
text_seqs = text_seqs[:, :-1]
if text_padding_mask is not None:
text_padding_mask = text_padding_mask.trim(1)
# Use the output of the text generator to compute the decoder output.
decoder_output, decoder_padding_mask = self.model.decode(
text_seqs,
text_padding_mask,
text_gen_output.encoder_output,
text_gen_output.encoder_padding_mask,
)
assert self.model.t2u_model is not None
assert self.unit_decoder is not None
unit_gen_output = None
prosody_encoder_out = None
if self.model.prosody_encoder_model is not None:
assert prosody_encoder_input is not None
prosody_input_seqs, prosody_padding_mask = get_seqs_and_padding_mask(
prosody_encoder_input
)
prosody_encoder_out = self.model.prosody_encoder_model(
prosody_input_seqs,
prosody_padding_mask,
).unsqueeze(1)
if isinstance(self.model.t2u_model, UnitYT2UModel):
assert self.unit_generator is not None
assert self.unit_prefix_indices is not None
# (S_pre) -> (N, S_pre)
prefix_seqs = self.unit_prefix_indices.expand(decoder_output.size(0), -1)
unit_gen_output = self.unit_generator(
source_seqs=decoder_output,
source_padding_mask=decoder_padding_mask,
prompt_seqs=prefix_seqs,
prompt_padding_mask=None,
)
assert self.model.t2u_model.target_vocab_info.pad_idx is not None
unit_seq_list = [h[0].seq for h in unit_gen_output.hypotheses]
unit_seqs, _ = pad_seqs(
unit_seq_list, self.model.t2u_model.target_vocab_info.pad_idx
)
else:
t2u_model_output, decoder_padding_mask, _ = self.model.t2u_model(
text_decoder_output=decoder_output,
text_decoder_padding_mask=decoder_padding_mask,
text_seqs=text_seqs,
duration_factor=duration_factor,
film_cond_emb=prosody_encoder_out,
)
# (B, S_unit, V_unit)
unit_seqs = t2u_model_output.logits.argmax(dim=2)
# Apply the padding mask to the generated units.
unit_seqs = apply_padding_mask(
unit_seqs, decoder_padding_mask, t2u_model_output.vocab_info.pad_idx
)
# Convert to speech units.
units = self.unit_decoder(unit_seqs)
# ngram-filtering doesn't apply to NAR unit decoding.
if ngram_filtering and isinstance(self.model.t2u_model, UnitYT2UModel):
if units.size(0) > 1:
raise NotImplementedError(
"unit ngram_filtering is not implemented for batch_size > 1."
)
arr = remove_consecutive_repeated_ngrams(units[0].tolist())
units = torch.tensor(arr).to(units).unsqueeze(0)
return texts, units
|