(self, encoder_dim, decoder_dim, num_heads, vocab_size, pad_id=-1)
| 1530 | |
| 1531 | class LASDecoderModelTest(TestCase): |
| 1532 | def _set_up_las(self, encoder_dim, decoder_dim, num_heads, vocab_size, pad_id=-1): |
| 1533 | num_layers = 2 |
| 1534 | cfg = LASDecoderModel.default_config().set( |
| 1535 | name="las_decoder", |
| 1536 | input_dim=encoder_dim, |
| 1537 | vocab_size=vocab_size, |
| 1538 | ) |
| 1539 | cfg.decoder = causal_lm.gpt_decoder_config( |
| 1540 | stack_cfg=attention.StackedTransformerLayer.default_config(), |
| 1541 | num_layers=num_layers, |
| 1542 | hidden_dim=decoder_dim, |
| 1543 | num_heads=num_heads, |
| 1544 | vocab_size=vocab_size, |
| 1545 | activation_function="nn.gelu", |
| 1546 | max_position_embeddings=10, |
| 1547 | layer_norm_epsilon=0.1, |
| 1548 | dropout_rate=0.0, |
| 1549 | ) |
| 1550 | cfg.decoder.pad_token_id = pad_id |
| 1551 | transformer_cfg = cfg.decoder.transformer.layer |
| 1552 | transformer_cfg.cross_attention = attention.TransformerAttentionLayer.default_config() |
| 1553 | transformer_cfg.cross_attention.attention.num_heads = num_heads |
| 1554 | return cfg |
| 1555 | |
| 1556 | def test_forward(self): |
| 1557 | encoder_dim, decoder_dim, num_heads, vocab_size = 8, 6, 3, 20 |
no test coverage detected