Merge branch 'add-tokenizer-save' into feat-model

Former-commit-id: 2cfaef4a020f43332a8f33b6a9bd8221ec9fae34

seq2seq/run_seq2seq_flax.py

@@ -818,13 +818,16 @@ def main():
                 params=params,
             )
 
+            # save tokenizer
+            tokenizer.save_pretrained(training_args.output_dir)
+
             # save state
             state = unreplicate(state)
             with (Path(training_args.output_dir) /  'opt_state.msgpack').open('wb') as f:
                 f.write(to_bytes(state.opt_state))
             with (Path(training_args.output_dir) /  'training_state.json').open('w') as f:
                 json.dump({'step': state.step.item()}, f)
-            
+
             # save to W&B
             if data_args.log_model:
                 metadata = {'step': step, 'epoch': epoch}
@@ -834,6 +837,11 @@ def main():
                     name=f"model-{wandb.run.id}", type="bart_model", metadata=metadata
                 )
                 artifact.add_file(str(Path(training_args.output_dir) / 'flax_model.msgpack'))
+                artifact.add_file(str(Path(training_args.output_dir) / 'tokenizer_config.json'))
+                artifact.add_file(str(Path(training_args.output_dir) / 'special_tokens_map.json'))
+                artifact.add_file(str(Path(training_args.output_dir) / 'vocab.json'))
+                artifact.add_file(str(Path(training_args.output_dir) / 'added_tokens.json'))
+                artifact.add_file(str(Path(training_args.output_dir) / 'merges.txt'))
                 artifact.add_file(str(Path(training_args.output_dir) / 'config.json'))
                 artifact.add_file(str(Path(training_args.output_dir) / 'opt_state.msgpack'))
                 artifact.add_file(str(Path(training_args.output_dir) / 'training_state.json'))