Repository files navigation

A follow up repository ofJax-Journey. This repository provides a selection of notebooks for various NLP tasks, which are completely see-through (i.e., you can see the implementation till the basic Jax/Haiku modules, in a single notebook). These were meant to be used as further tutorials in Jax for NLP, and as a guide for the coding style followed in thisawesome article by Madison May.

These notebooks, although mostly code, also mention the nuanced features, often missed when using off-the-shelf models. Moreover, they allow you to optimize everything right to the innermost modules. Also, we mention how to adapt the model to your use case, in each notebook.

A basic introductory notebook consisting of the originalRoBERTa initialized version andrandomly initialized version .

Here we realise the need for restructuring the code, and correspondingly, place all the code component-wise insrc/ . The new things we code over the original implementation are:

• The masking function for MLMhere,
• AHuggingFace Tokenizers based tokenizer,here
• A Language Embedding for TLM task,here.
• Additionally, we include an option to make the transformer auto-regressive and add a mask for the same,here. This is needed for CLM.

The final notebook can be foundhere.

Hoping to create a kind, giving, an open community that forms deep connetions by working together. Join us here :https://discord.gg/s6xSHG94u5