Type annotate transformer/modules.py

[spencerp]

Patch description
Problem:

• When using these modules, it's not immediately obvious what types of args they expect and return. You need to read through the implementation or add breakpoints to figure this out, which wastes time.

Solution:

• Type annotate __init__ and forward functions that aren't already annotated.

There shouldn't be any functional changes here, these are essentially just comments with very loose static checking.

Testing steps
Ran a sample transformer training run to sanity check:

λ parlai train_model --task twitter --model-file /tmp/tr_twitter --model transformer/ranker --batchsize 16 --validation-every-n-secs 3600 --candidates batch --eval-candidates batch --data-parallel True

Circle CI.

[spencerp]

I had to move these down so TransformerEncoder is declared before using it as a type. If we want to avoid moving these, I could just wrap the type in quotes "TransformerEncoder", but that disables static checking of these.

[stephenroller]

Ah the secret to that is to add at the top:

from __future__ import annotations

This causes python to delay typechecking

[spencerp]

This ended up being resolved anyway by the breaking up of modules.py

[stephenroller]

thx! suggesting we do the future annotations.