Add bundles for FLAN-T5 (#2061)

* Add initial bundles for FLAN-T5 * Fix path * Add docstrings * Fix transform error * Trim whitespace * Add modules to init
pytorch · Feb 17, 2023 · 21261ea · 21261ea
1 parent 19f8bc9
commit 21261ea
Show file tree

Hide file tree

Showing 2 changed files with 352 additions and 92 deletions.
diff --git a/torchtext/prototype/models/t5/__init__.py b/torchtext/prototype/models/t5/__init__.py
@@ -1,4 +1,19 @@
 from .bundler import (
+    FLAN_T5_SMALL_ENCODER,
+    FLAN_T5_SMALL,
+    FLAN_T5_SMALL_GENERATION,
+    FLAN_T5_BASE_ENCODER,
+    FLAN_T5_BASE,
+    FLAN_T5_BASE_GENERATION,
+    FLAN_T5_LARGE_ENCODER,
+    FLAN_T5_LARGE,
+    FLAN_T5_LARGE_GENERATION,
+    FLAN_T5_XL_ENCODER,
+    FLAN_T5_XL,
+    FLAN_T5_XL_GENERATION,
+    FLAN_T5_XXL_ENCODER,
+    FLAN_T5_XXL,
+    FLAN_T5_XXL_GENERATION,
     T5_11B,
     T5_11B_ENCODER,
     T5_11B_GENERATION,
@@ -38,5 +53,20 @@
     "T5_11B_ENCODER",
     "T5_11B",
     "T5_11B_GENERATION",
+    "FLAN_T5_SMALL_ENCODER",
+    "FLAN_T5_SMALL",
+    "FLAN_T5_SMALL_GENERATION",
+    "FLAN_T5_BASE_ENCODER",
+    "FLAN_T5_BASE",
+    "FLAN_T5_BASE_GENERATION",
+    "FLAN_T5_LARGE_ENCODER",
+    "FLAN_T5_LARGE",
+    "FLAN_T5_LARGE_GENERATION",
+    "FLAN_T5_XL_ENCODER",
+    "FLAN_T5_XL",
+    "FLAN_T5_XL_GENERATION",
+    "FLAN_T5_XXL_ENCODER",
+    "FLAN_T5_XXL",
+    "FLAN_T5_XXL_GENERATION",
     "T5Transform",
 ]