Update extract_feature_print.py

extract_feature_print.py

@@ -157,9 +157,9 @@ def main():
 
     # Resolve model path and name
     custom_mappings = {
+        "wav2vec_2" : ("wav2vec_small_960h.pt", "wav2vec"),
         "hubert_base": ("hubert_base.pt", "hubert"),
         "contentvec_base": ("contentvec_base.pt", "contentvec"),
-        "hubert_large_ll60k": ("hubert_large_ll60k.pt", "hubert"),
     }
     if os.path.split(model_path)[-1] == "Custom" and model_name in custom_mappings:
         model_path, resolved_model_name = custom_mappings[model_name]
@@ -170,7 +170,7 @@ def main():
         sys.exit(1)
 
     # Load model
-    model_config = model_configs.get(model_name, model_configs["hubert"])
+    model_config = model_configs.get(model_name, model_configs[resolved_model_name])
     model_dict = model_config["load_model"](model_path, config.device, config.is_half)
     model = model_dict["model"]
     additional_configs = model_dict.get("saved_cfg")