Added configs.

mlfoundations · Oct 11, 2023 · 902f816 · 902f816
1 parent 4ccb752
commit 902f816
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 0 deletions.
diff --git a/src/open_clip/hf_configs.py b/src/open_clip/hf_configs.py
@@ -53,4 +53,15 @@
         },
         "pooler": "cls_pooler",
     },
+    # https://huggingface.co/docs/transformers/model_doc/m2m_100
+    "m2m_100": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "d_model",
+            "heads": "encoder_attention_heads",
+            "layers": "encoder_layers",
+        },
+        "pooler": "cls_pooler",
+    },
 }
diff --git a/src/open_clip/model_configs/nllb-clip-base.json b/src/open_clip/model_configs/nllb-clip-base.json
@@ -0,0 +1,15 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "hf_model_name": "facebook/nllb-200-distilled-600M",
+        "hf_tokenizer_name": "facebook/nllb-200-distilled-600M",
+        "proj": "linear",
+        "pooler_type": "cls_pooler"
+    }
+}
diff --git a/src/open_clip/model_configs/nllb-clip-large.json b/src/open_clip/model_configs/nllb-clip-large.json
@@ -0,0 +1,16 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 32,
+        "width": 1280,
+        "head_width": 80,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "hf_model_name": "facebook/nllb-200-distilled-1.3B",
+        "hf_tokenizer_name": "facebook/nllb-200-distilled-1.3B",
+        "proj": "linear",
+        "pooler_type": "cls_pooler"
+    }
+}