(
self,
vocab_size: int,
encoder_output_size: int,
attention_heads: int = 4,
linear_units: int = 2048,
num_blocks: int = 6,
dropout_rate: float = 0.1,
positional_dropout_rate: float = 0.1,
self_attention_dropout_rate: float = 0.0,
src_attention_dropout_rate: float = 0.0,
input_layer: str = "embed",
use_output_layer: bool = True,
normalize_before: bool = True,
src_attention: bool = True,
key_bias: bool = True,
activation_type: str = "relu",
gradient_checkpointing: bool = False,
tie_word_embedding: bool = False,
)
| 56 | """ |
| 57 | |
| 58 | def __init__( |
| 59 | self, |
| 60 | vocab_size: int, |
| 61 | encoder_output_size: int, |
| 62 | attention_heads: int = 4, |
| 63 | linear_units: int = 2048, |
| 64 | num_blocks: int = 6, |
| 65 | dropout_rate: float = 0.1, |
| 66 | positional_dropout_rate: float = 0.1, |
| 67 | self_attention_dropout_rate: float = 0.0, |
| 68 | src_attention_dropout_rate: float = 0.0, |
| 69 | input_layer: str = "embed", |
| 70 | use_output_layer: bool = True, |
| 71 | normalize_before: bool = True, |
| 72 | src_attention: bool = True, |
| 73 | key_bias: bool = True, |
| 74 | activation_type: str = "relu", |
| 75 | gradient_checkpointing: bool = False, |
| 76 | tie_word_embedding: bool = False, |
| 77 | ): |
| 78 | super().__init__() |
| 79 | attention_dim = encoder_output_size |
| 80 | activation = INSPIREMUSIC_ACTIVATION_CLASSES[activation_type]() |
| 81 | |
| 82 | self.embed = torch.nn.Sequential( |
| 83 | torch.nn.Identity() if input_layer == "no_pos" else |
| 84 | torch.nn.Embedding(vocab_size, attention_dim), |
| 85 | INSPIREMUSIC_EMB_CLASSES[input_layer](attention_dim, |
| 86 | positional_dropout_rate), |
| 87 | ) |
| 88 | |
| 89 | self.normalize_before = normalize_before |
| 90 | self.after_norm = torch.nn.LayerNorm(attention_dim, eps=1e-5) |
| 91 | self.use_output_layer = use_output_layer |
| 92 | if use_output_layer: |
| 93 | self.output_layer = torch.nn.Linear(attention_dim, vocab_size) |
| 94 | else: |
| 95 | self.output_layer = torch.nn.Identity() |
| 96 | self.num_blocks = num_blocks |
| 97 | self.decoders = torch.nn.ModuleList([ |
| 98 | DecoderLayer( |
| 99 | attention_dim, |
| 100 | INSPIREMUSIC_ATTENTION_CLASSES["selfattn"]( |
| 101 | attention_heads, attention_dim, |
| 102 | self_attention_dropout_rate, key_bias), |
| 103 | INSPIREMUSIC_ATTENTION_CLASSES["selfattn"]( |
| 104 | attention_heads, attention_dim, src_attention_dropout_rate, |
| 105 | key_bias) if src_attention else None, |
| 106 | PositionwiseFeedForward(attention_dim, linear_units, |
| 107 | dropout_rate, activation), |
| 108 | dropout_rate, |
| 109 | normalize_before, |
| 110 | ) for _ in range(self.num_blocks) |
| 111 | ]) |
| 112 | |
| 113 | self.gradient_checkpointing = gradient_checkpointing |
| 114 | self.tie_word_embedding = tie_word_embedding |
| 115 |
no test coverage detected