minor missing key fixes

mosaicml · Apr 13, 2024 · f29191d · f29191d
1 parent 6659531
commit f29191d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 289 deletions.
diff --git a/convert_mosaicbert_to_hf.py b/convert_mosaicbert_to_hf.py
diff --git a/scripts/inference/convert_mosaicbert_to_hf.py b/scripts/inference/convert_mosaicbert_to_hf.py
@@ -82,11 +82,11 @@ def write_huggingface_pretrained_from_composer_checkpoint(
         #json
         with open(bert_config_path, 'r') as f:
             bert_config = json.load(f)   
-        composer_state_dict["state"]["integrations"]["huggingface"]["model"]["config"]["content"] = bert_config
+        composer_state_dict["state"]["integrations"] = {"huggingface":{"model":{"config":{"content": bert_config}}}}
 
     else:
         # placeholder config from mosaicml/mosaic-bert-base
-        composer_state_dict["state"]["integrations"]["huggingface"] = {
+        composer_state_dict["state"]["integrations"]={"huggingface":{
             "model":{"config":{"content":{
                 "_name_or_path": "mosaicml/mosaic-bert",
                 "alibi_starting_size": 512,
@@ -118,7 +118,7 @@ def write_huggingface_pretrained_from_composer_checkpoint(
                 "use_cache": False,
                 "vocab_size": 30522
 
-        }}},
+        }}}},
         "tokenizer":{}
     }
     if 'state' not in composer_state_dict: