Updating examples/nlp_seq

[andsteing]

This issue tracks updates the nlp_seq example to follow practices outlined in #231.

• [ ] Port to linen API - once ported all subsequent changes should be done in linen_examples/nlp_seq
• [ ] Update README.md (links to paper, command line, TensorBoard metrics)
• [ ] Add requirements.txt
• [ ] Update file structure (separate main.py, configs/ directory)
• [ ] Use ml_collections.ConfigDict
• [ ] Add benchmark test (optional)
• [ ] Add unit test
• [ ] Add Colab (optional)
• [ ] Adhere to Google Python style
• [ ] Shorten/beautify training loop (consider using clu for this)

[andsteing]

Note that "file structure" in #231 got changed (as part of #634):

• To make it easier to test, maintain and reuse code structure the code as follows:

  • main.py contains the flags and calls a method from train.py to run the training loop. This should be the only file defining flags! This can be almost identical for all examples, please copy from linen_examples/wmt/main.py.
  • train.py contains classes and methods for training and evaluating the model.
  • train_test.py for test cases of the training code. At a minimum the test should run a single training step but more fine grained unit tests are a bonus. You can use tfds.testing.mock_data to avoid real data from disk.

[rozeappletree]

Anyone working on this? I can take it up.

[bohnetbd]

I commited to update it and promised to do it. However, I am on holiday for the next two weeks. What would you like to do? The config needs to be replaced and the optimizer updated.

[rozeappletree]

I commited to update it and promised to do it. However, I am on holiday for the next two weeks. What would you like to do? The config needs to be replaced and the optimizer updated.

Oh. I was looking to learn flax by sending some PRs. This issue seemed perfect to familiarise myself with flax way of doing things.

I can pick up the two tasks above. May need some help at some places.

[bohnetbd]

Cool. Sounds good to me. Personally, I would recommend to take on one task at a time :) so it is easier to review.

[andsteing]

One task at a time LGTM.

When doing non-trivial changes (e.g. updating to Optax), thanks for adding links to before/after metrics like in #1476