more resilient

mosaicml · May 27, 2024 · d79ba7c · d79ba7c
1 parent 76921d8
commit d79ba7c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llmfoundry/utils/huggingface_hub_utils.py b/llmfoundry/utils/huggingface_hub_utils.py
@@ -240,10 +240,10 @@ def edit_files_for_hf_compatibility(
     with open(os.path.join(folder, 'config.json'), 'r') as _f:
         config = json.load(_f)
 
-    # If the config file exists, the entrypoint files are specified in the auto map
+    # If the config file exists, the entrypoint files would be specified in the auto map
     entrypoint_files = set()
     if config_file_exists:
-        for key, value in config['auto_map'].items():
+        for key, value in config.get('auto_map', {}).items():
             # Only keep the modeling entrypoints, e.g. AutoModelForCausalLM
             if 'model' not in key.lower():
                 continue