Add mt5 to the mapping of flex heads

adapter-hub · Jan 5, 2024 · 39769cf · 39769cf
1 parent 19f4506
commit 39769cf
Showing 1 changed file with 29 additions and 0 deletions.
diff --git a/src/adapters/head_utils.py b/src/adapters/head_utils.py
@@ -369,6 +369,35 @@
         },
         "layers": ["lm_head"],
     },
+    # MT5
+        "MT5ForConditionalGeneration": {
+        "config": {
+            "head_type": "seq2seq_lm",
+        },
+        "layers": ["lm_head"],
+    },
+    "MT5ForQuestionAnswering": {
+        "config": {
+            "head_type": "question_answering",
+            "activation_function": None,
+            "layers": 1,
+        },
+        "layers": [None, "qa_outputs"],
+    },
+    "MT5ForSequenceClassification": {
+        "config": {
+            "head_type": "classification",
+            "layers": 2,
+            "activation_function": "tanh",
+        },
+        "layers": [
+            None,
+            "classification_head.dense",
+            None,
+            None,
+            "classification_head.out_proj",
+        ],
+    },
     # DistilBERT
     "DistilBertForSequenceClassification": {
         "config": {