Add Encoder-Decoder model support - #30
Conversation
| with open(config_path, "r") as f: | ||
| config = yaml.safe_load(f) |
There was a problem hiding this comment.
I think we have helper method for yaml dumping and loading in config.py
|
|
||
| with open(config_path, "w") as f: | ||
| yaml.safe_dump(config, f) |
|
|
||
| # Dynamic routing for encoder-decoder setup | ||
| dataset_implementation_path = dataset_spec.implementation_path | ||
| trainer_implementation_path = trainer_spec.implementation_path | ||
| dataset_readme_path = dataset_spec.readme_template_path | ||
| trainer_readme_path = trainer_spec.readme_template_path | ||
|
|
||
| if model_name == "encoder-decoder": | ||
| if dataset_name == "string-reverse": | ||
| dataset_implementation_path = Path( | ||
| "trainite/datasets/string_reverse_seq2seq.py" | ||
| ) | ||
| dataset_readme_path = Path( | ||
| "trainite/templates/components/datasets/string_reverse_seq2seq.md" | ||
| ) | ||
| if trainer_name == "pretrainer": | ||
| trainer_implementation_path = Path( | ||
| "trainite/trainers/pretrainer_seq2seq.py" | ||
| ) | ||
| trainer_readme_path = Path( | ||
| "trainite/templates/components/trainers/pretrainer_seq2seq.md" | ||
| ) | ||
|
|
There was a problem hiding this comment.
Dont understand use of it, I thought dataset, trainer and model would be independent you might have to consult with victor, on how modular should the generated code be, can user select different components or based on the use case.
| self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) | ||
| self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False) |
There was a problem hiding this comment.
we can have a single multi dim linear layer for key and value as it would be both either x in case of self and y in case of cross.
| self.encoder_pos_encoding = PositionalEncoding( | ||
| hidden_size, encoder_max_seq_len, dropout | ||
| ) | ||
| self.decoder_pos_encoding = PositionalEncoding( | ||
| hidden_size, decoder_max_seq_len, dropout | ||
| ) |
There was a problem hiding this comment.
we could use single layer for this
There was a problem hiding this comment.
I wanted to keep them decoupled incase someone wants to try and modify either one of encoder or decoder in the setup.
There was a problem hiding this comment.
I think the model should be simple enough, if your wants separate encoder , decoder position encoder they can add it themselves
| if (encoder_input_ids == self.embedding.padding_idx).any(): | ||
| cross_padding_mask = ( | ||
| encoder_input_ids != self.embedding.padding_idx | ||
| ).reshape(B_dec, 1, 1, S_enc) | ||
| else: | ||
| cross_padding_mask = None |
There was a problem hiding this comment.
its same as encoder self padding?
|
Also like I dont think we need new datasets for this model old reverse-string should work and the old trainer should work too. |
I would discuss this with victor as well, but having the same dataset and pretrainer for both cases can lead to if/else statements inside the file which can be more confusing for the end-user. That is why I opted for creating a seperate file to keep it clean and only the code required for the model they chose. |
29b6131 to
25ef6d9
Compare
2 participants
PR Description:
This Pull Request introduces the Encoder-Decoder (Seq2Seq) transformer architecture option to the Trainite CLI toolbox. It fully decouples encoder and decoder hyperparameters, implements robust sequence-to-sequence template routing, and establishes a highly thorough, symmetrical testing suite across all datasets, models, and trainers.
1. Symmetrical & Decoupled Model Architecture
trainite/models/encoder_decoder.pywith:MultiHeadAttentionsupporting causal self-attention, non-causal self-attention, and cross-attention.EncoderBlockandDecoderBlocklayer stacks.EncoderDecoderModelConfig(num_encoder_layersvsnum_decoder_layers,encoder_max_seq_lenvsdecoder_max_seq_len).2. Dynamic Template Routing (CLI)
transformertodecoderfor symmetric clarity.--model encoder-decoderoption in the CLI (trainite/cli/init.py).if/elsechecks or branching logic into user-generated boilerplates.3. Specialized Seq2Seq Dataset & Trainer Templates
trainite/datasets/string_reverse_seq2seq.py): Yields separate source and target tensors with bespoke padding collation (handling target pads withignore_index=-100and inputs with0).trainite/trainers/pretrainer_seq2seq.py): Implements clean pretraining and evaluation loops specifically for sequence-to-sequence forward passes.4. Fully Aligned and Comprehensive Testing
tests/models/encoder_decoder_test.py): Added positional encoding dimension checks, spec-based model instantiation, and deterministic/random dropout evaluations.tests/datasets/string_reverse_seq2seq_test.py): Added checks for special tokens, character preset variations, sequence lengths, and exception handling.tests/trainers/pretrainer_seq2seq_test.py): Integrated trainer lifecycle tests verifying loss flattening, exact accuracy conversions, device selections, and last-checkpoint fallbacks.