Allow flan-t5 models in ParlAI with fp16 improvment

parlai/agents/hugging_face/t5.py

@@ -41,8 +41,9 @@ def check_hf_version(v: Tuple[int, int]) -> bool:
 def build_t5(opt: Opt) -> T5ForConditionalGeneration:
     if not check_hf_version(HF_VERSION):
         raise RuntimeError('Must use transformers package >= 4.3 to use t5')
+    torch_dtype = torch.float16 if opt['fp16'] else torch.float32
     return T5ForConditionalGeneration.from_pretrained(
-        opt['t5_model_arch'], dropout_rate=opt['t5_dropout']
+        opt['t5_model_arch'], dropout_rate=opt['t5_dropout'], torch_dtype=torch_dtype
     )
 
 
@@ -86,7 +87,18 @@ def add_cmdline_args(
             '--t5-model-arch',
             type=str,
             default='t5-base',
-            choices=["t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b"],
+            choices=[
+                "t5-small",
+                "t5-base",
+                "t5-large",
+                "t5-3b",
+                "t5-11b",
+                "google/flan-t5-small",
+                "google/flan-t5-base",
+                "google/flan-t5-large",
+                "google/flan-t5-xl",
+                "google/flan-t5-xxl",
+            ],
         )
         group.add_argument(
             '--t5-model-parallel',