export_model tweak (#1277)

* export_model tweak

* add hf_pretrained_model_name_or_path to saved config

* remove model_type form export_model

* update test

* remote unused var

Author: zphang

examples/notebooks/jiant_Basic_Example.ipynb

@@ -310,8 +310,8 @@
     "    jiant_task_container_config_path=\"./run_configs/mrpc_run_config.json\",\n",
     "    output_dir=\"./runs/mrpc\",\n",
     "    hf_pretrained_model_name_or_path=\"roberta-base\",\n",
-    "    model_path=\"./models/roberta-base/model/roberta.p\",\n",
-    "    model_config_path=\"./models/roberta-base/model/roberta.json\",\n",
+    "    model_path=\"./models/roberta-base/model/model.p\",\n",
+    "    model_config_path=\"./models/roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",

examples/notebooks/jiant_EdgeProbing_Example.ipynb

@@ -3147,8 +3147,8 @@
     "    jiant_task_container_config_path=\"./run_configs/semeval_run_config.json\",\n",
     "    output_dir=\"./runs/semeval\",\n",
     "    hf_pretrained_model_name_or_path=\"roberta-base\",\n",
-    "    model_path=\"./models/roberta-base/model/roberta.p\",\n",
-    "    model_config_path=\"./models/roberta-base/model/roberta.json\",\n",
+    "    model_path=\"./models/roberta-base/model/model.p\",\n",
+    "    model_config_path=\"./models/roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",

examples/notebooks/jiant_MNLI_Diagnostic_Example.ipynb

@@ -322,8 +322,8 @@
     "    jiant_task_container_config_path=\"./run_configs/jiant_run_config.json\",\n",
     "    output_dir=\"./runs/run1\",\n",
     "    hf_pretrained_model_name_or_path=\"roberta-base\",\n",
-    "    model_path=\"./models/roberta-base/model/roberta.p\",\n",
-    "    model_config_path=\"./models/roberta-base/model/roberta.json\",\n",
+    "    model_path=\"./models/roberta-base/model/model.p\",\n",
+    "    model_config_path=\"./models/roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",

examples/notebooks/jiant_Multi_Task_Example.ipynb

@@ -343,8 +343,8 @@
     "    jiant_task_container_config_path=\"./run_configs/jiant_run_config.json\",\n",
     "    output_dir=\"./runs/run1\",\n",
     "    hf_pretrained_model_name_or_path=\"roberta-base\",\n",
-    "    model_path=\"./models/roberta-base/model/roberta.p\",\n",
-    "    model_config_path=\"./models/roberta-base/model/roberta.json\",\n",
+    "    model_path=\"./models/roberta-base/model/model.p\",\n",
+    "    model_config_path=\"./models/roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",

examples/notebooks/jiant_STILTs_Example.ipynb

@@ -368,8 +368,8 @@
     "    jiant_task_container_config_path=\"./run_configs/mnli_run_config.json\",\n",
     "    output_dir=\"./runs/mnli\",\n",
     "    hf_pretrained_model_name_or_path=\"roberta-base\",\n",
-    "    model_path=\"./models/roberta-base/model/roberta.p\",\n",
-    "    model_config_path=\"./models/roberta-base/model/roberta.json\",\n",
+    "    model_path=\"./models/roberta-base/model/modle.p\",\n",
+    "    model_config_path=\"./models/roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",
@@ -406,7 +406,7 @@
     "    hf_pretrained_model_name_or_path=\"roberta-base\",\n",
     "    model_path=\"./runs/mnli/best_model.p\",  # Loading the best model\n",
     "    model_load_mode=\"partial\",\n",
-    "    model_config_path=\"./models/roberta-base/model/roberta.json\",\n",
+    "    model_config_path=\"./models/roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",

examples/notebooks/jiant_XNLI_Example.ipynb

@@ -384,8 +384,8 @@
     "    jiant_task_container_config_path=\"./run_configs/jiant_run_config.json\",\n",
     "    output_dir=\"./runs/run1\",\n",
     "    hf_pretrained_model_name_or_path=\"xlm-roberta-base\",\n",
-    "    model_path=\"./models/xlm-roberta-base/model/xlm-roberta.p\",\n",
-    "    model_config_path=\"./models/xlm-roberta-base/model/xlm-roberta.json\",\n",
+    "    model_path=\"./models/xlm-roberta-base/model/model.p\",\n",
+    "    model_config_path=\"./models/xlm-roberta-base/model/config.json\",\n",
     "    learning_rate=1e-5,\n",
     "    eval_every_steps=500,\n",
     "    do_train=True,\n",

jiant/proj/main/export_model.py

@@ -41,11 +41,10 @@ def export_model(
         output_base_path: Base path to save output to
     """
     model = AutoModelForPreTraining.from_pretrained(hf_pretrained_model_name_or_path)
-    model_type = model.config_class.model_type
 
     model_fol_path = os.path.join(output_base_path, "model")
-    model_path = os.path.join(model_fol_path, f"{model_type}.p")
-    model_config_path = os.path.join(model_fol_path, f"{model_type}.json")
+    model_path = os.path.join(model_fol_path, "model.p")
+    model_config_path = os.path.join(model_fol_path, "config.json")
     tokenizer_fol_path = os.path.join(output_base_path, "tokenizer")
 
     os.makedirs(tokenizer_fol_path, exist_ok=True)
@@ -56,7 +55,7 @@ def export_model(
     tokenizer = AutoTokenizer.from_pretrained(hf_pretrained_model_name_or_path)
     tokenizer.save_pretrained(tokenizer_fol_path)
     config = {
-        "model_type": model_type,
+        "hf_pretrained_model_name_or_path": hf_pretrained_model_name_or_path,
         "model_path": model_path,
         "model_config_path": model_config_path,
     }

tests/proj/main/test_export_model.py

@@ -19,6 +19,6 @@ def test_export_model(tmp_path, model_type, model_class, hf_pretrained_model_nam
         output_base_path=tmp_path,
     )
     read_config = py_io.read_json(os.path.join(tmp_path, f"config.json"))
-    assert read_config["model_type"] == model_type
-    assert read_config["model_path"] == os.path.join(tmp_path, "model", f"{model_type}.p")
-    assert read_config["model_config_path"] == os.path.join(tmp_path, "model", f"{model_type}.json")
+    assert read_config["hf_pretrained_model_name_or_path"] == hf_pretrained_model_name_or_path
+    assert read_config["model_path"] == os.path.join(tmp_path, "model", "model.p")
+    assert read_config["model_config_path"] == os.path.join(tmp_path, "model", "config.json")