Add missing ESM autoclass (huggingface#20177)

* Add missing ESM autoclass * Correct ESMFold checkpoint
altryne · Nov 15, 2022 · d3d5fa3 · d3d5fa3
1 parent 92cfe8b
commit d3d5fa3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py
@@ -445,6 +445,7 @@
         ("distilbert", "DistilBertForMaskedLM"),
         ("electra", "ElectraForMaskedLM"),
         ("ernie", "ErnieForMaskedLM"),
+        ("esm", "EsmForMaskedLM"),
         ("flaubert", "FlaubertWithLMHeadModel"),
         ("fnet", "FNetForMaskedLM"),
         ("funnel", "FunnelForMaskedLM"),

diff --git a/src/transformers/models/esm/modeling_esmfold.py b/src/transformers/models/esm/modeling_esmfold.py
@@ -52,7 +52,7 @@
 
 
 logger = logging.get_logger(__name__)
-_CHECKPOINT_FOR_DOC = "Rocketknight1/esmfold_v1"
+_CHECKPOINT_FOR_DOC = "facebook/esmfold_v1"
 _CONFIG_FOR_DOC = "EsmConfig"
 _TOKENIZER_FOR_DOC = "EsmTokenizer"