Fix normalized config key for models architecture

huggingface · Sep 22, 2023 · d6f8e10 · d6f8e10
1 parent 6d1ae0e
commit d6f8e10
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -209,7 +209,7 @@ class NormalizedConfigManager:
         # "big_bird": NormalizedTextConfig,
         # "bigbird_pegasus": BartLikeNormalizedTextConfig,
         "blenderbot": BartLikeNormalizedTextConfig,
-        "blenderbot_small": BartLikeNormalizedTextConfig,
+        "blenderbot-small": BartLikeNormalizedTextConfig,
         "bloom": NormalizedTextConfig.with_args(num_layers="n_layer"),
         "falcon": NormalizedTextConfig.with_args(num_layers="num_hidden_layers", num_attention_heads="num_kv_heads"),
         "camembert": NormalizedTextConfig,
@@ -224,16 +224,16 @@ class NormalizedConfigManager:
         "encoder-decoder": NormalizedEncoderDecoderConfig,
         "gpt2": GPT2LikeNormalizedTextConfig,
         "gpt-bigcode": GPT2LikeNormalizedTextConfig,
-        "gpt_neo": NormalizedTextConfig.with_args(num_attention_heads="num_heads"),
-        "gpt_neox": NormalizedTextConfig,
+        "gpt-neo": NormalizedTextConfig.with_args(num_attention_heads="num_heads"),
+        "gpt-neox": NormalizedTextConfig,
         "llama": NormalizedTextConfig,
         "gptj": GPT2LikeNormalizedTextConfig,
         "imagegpt": GPT2LikeNormalizedTextConfig,
         "longt5": T5LikeNormalizedTextConfig,
         "marian": BartLikeNormalizedTextConfig,
         "mbart": BartLikeNormalizedTextConfig,
         "mt5": T5LikeNormalizedTextConfig,
-        "m2m_100": BartLikeNormalizedTextConfig,
+        "m2m-100": BartLikeNormalizedTextConfig,
         "nystromformer": NormalizedTextConfig,
         "opt": NormalizedTextConfig,
         "pegasus": BartLikeNormalizedTextConfig,
@@ -242,7 +242,7 @@ class NormalizedConfigManager:
         "regnet": NormalizedVisionConfig,
         "resnet": NormalizedVisionConfig,
         "roberta": NormalizedTextConfig,
-        "speech_to_text": SpeechToTextLikeNormalizedTextConfig,
+        "speech-to-text": SpeechToTextLikeNormalizedTextConfig,
         "splinter": NormalizedTextConfig,
         "t5": T5LikeNormalizedTextConfig,
         "trocr": TrOCRLikeNormalizedTextConfig,
@@ -252,7 +252,7 @@ class NormalizedConfigManager:
         "xlm-roberta": NormalizedTextConfig,
         "yolos": NormalizedVisionConfig,
         "mpt": MPTNormalizedTextConfig,
-        "gpt_bigcode": GPTBigCodeNormalizedTextConfig,
+        "gpt-bigcode": GPTBigCodeNormalizedTextConfig,
     }
 
     @classmethod
@@ -266,5 +266,6 @@ def check_supported_model(cls, model_type: str):
 
     @classmethod
     def get_normalized_config_class(cls, model_type: str) -> Type:
+        model_type = model_type.replace("_", "-")
         cls.check_supported_model(model_type)
         return cls._conf[model_type]