| 114 | |
| 115 | |
| 116 | class EncoderDecoderConfig: |
| 117 | def __init__(self, d_model, vocab_size, num_heads, d_kv, d_ff, num_layers, |
| 118 | relative_attention_num_buckets_or_max_pos_seq_len, decoder_start_token_id=0, decoder_end_token_id=1): |
| 119 | self.d_model = d_model |
| 120 | self.vocab_size = vocab_size |
| 121 | self.num_heads = num_heads |
| 122 | self.d_kv = d_kv |
| 123 | self.d_ff = d_ff |
| 124 | self.num_layers = num_layers |
| 125 | self.relative_attention_num_buckets = relative_attention_num_buckets_or_max_pos_seq_len |
| 126 | self.decoder_start_token_id = decoder_start_token_id |
| 127 | self.decoder_end_token_id = decoder_end_token_id |
| 128 | |
| 129 | |
| 130 | data_type_mapping = {"fp32": 0, "fp16": 1, "bf16": 2} |