add wiki sentences

README.md

@@ -4,7 +4,7 @@ A Transformer-VAE made using flax.
 
 Done as part of Huggingface community training ([see forum post](https://discuss.huggingface.co/t/train-a-vae-to-interpolate-on-english-sentences/7548)).
 
-Builds on T5, using an autoencoder to convert it into a VAE.
+Builds on T5, using an autoencoder to convert it into an MMD-VAE.
 
 [See training logs.](https://wandb.ai/fraser/flax-vae)
 

train.py

@@ -2,7 +2,7 @@
     Pre-training/Fine-tuning seq2seq models on autoencoding a dataset.
 
     TODO:
-    - [ ] Get this running.
+    - [x] Get this running.
     - [x] Don't make decoder input ids.
     - [ ] Add reg loss
         - [x] calculate MMD loss
@@ -372,6 +372,13 @@ def main():
             config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
         )
 
+    if model_args.add_special_tokens:
+        special_tokens_dict = {'pad_token': '<PAD>', 'bos_token': '<BOS>', 'eos_token': '<EOS>'}
+        num_added_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+        print('We have added', num_added_tokens, 'tokens to GPT2')
+        model.resize_token_embeddings(len(tokenizer))
+        assert tokenizer.pad_token == '<PAD>'
+
     # Preprocessing the datasets.
     # First we tokenize all the texts.
     if training_args.do_train:

wiki_sentences.sh

@@ -0,0 +1,23 @@
+export RUN_NAME=wiki_split
+
+./venv/bin/python train.py \
+--t5_model_name_or_path="t5-base" \
+--tokenizer_name=gpt2 \
+--add_special_tokens \
+--output_dir="output/${RUN_NAME}" \
+--overwrite_output_dir \
+--dataset_name=./datasets/wiki_sentences \
+--input_ids_column=token_ids \
+--do_train --do_eval \
+--n_latent_tokens 6 \
+--latent_token_size 32 \
+--save_steps="2500" \
+--eval_steps="2500" \
+--block_size="256" \
+--per_device_train_batch_size="10" \
+--per_device_eval_batch_size="10" \
+--learning_rate="5e-3" --warmup_steps="1000" \
+--adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
+--overwrite_output_dir \
+--num_train_epochs="1" \
+--push_to_hub \

wiki_split.sh

@@ -0,0 +1,20 @@
+export RUN_NAME=wiki_split
+
+./venv/bin/python train.py \
+--t5_model_name_or_path="t5-base" \
+--output_dir="output/${RUN_NAME}" \
+--overwrite_output_dir \
+--dataset_name="wiki_split" \
+--do_train --do_eval \
+--n_latent_tokens 6 \
+--latent_token_size 32 \
+--save_steps="2500" \
+--eval_steps="2500" \
+--block_size="32" \
+--per_device_train_batch_size="10" \
+--per_device_eval_batch_size="10" \
+--learning_rate="5e-3" --warmup_steps="1000" \
+--adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
+--overwrite_output_dir \
+--num_train_epochs="3" \
+--push_to_hub \