File size: 2,034 Bytes
5ec225a |
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 |
# 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 typing import Optional, Tuple, final
from fairseq2.models.transformer.frontend import TransformerFrontend
from fairseq2.nn.incremental_state import IncrementalStateBag
from fairseq2.nn.padding import PaddingMask
from fairseq2.nn.projection import Projection
from overrides import final as finaloverride
from torch import Tensor
from torch.nn import Module
from seamless_communication.models.monotonic_decoder.monotonic_decoder import (
MonotonicTransformerDecoder,
)
@final
class MonotonicDecoderModel(Module):
text_decoder_frontend: TransformerFrontend
text_decoder: MonotonicTransformerDecoder
final_proj: Projection
def __init__(
self,
text_decoder_frontend: TransformerFrontend,
text_decoder: MonotonicTransformerDecoder,
final_proj: Projection,
) -> None:
super().__init__()
self.text_decoder_frontend = text_decoder_frontend
self.text_decoder = text_decoder
self.final_proj = final_proj
@finaloverride
def decode(
self,
seqs: Tensor,
padding_mask: Optional[PaddingMask],
encoder_output: Tensor,
encoder_padding_mask: Optional[PaddingMask],
*,
state_bag: Optional[IncrementalStateBag] = None,
) -> Tuple[Tensor, Optional[PaddingMask], Tensor]:
seqs, padding_mask = self.text_decoder_frontend(
seqs, padding_mask, state_bag=state_bag
)
return self.text_decoder( # type: ignore[no-any-return]
seqs,
padding_mask,
encoder_output,
encoder_padding_mask,
state_bag=state_bag,
)
@finaloverride
def project(self, decoder_output: Tensor) -> Tensor:
logits = self.final_proj(decoder_output)
return logits # type: ignore[no-any-return]
|