Merge pull request #416 from abhi1092/main

Using AutoConfig to load model config file.

src/instructlab/training/main_ds.py

@@ -4,7 +4,6 @@
 from copy import deepcopy
 from pathlib import Path
 import argparse
-import json
 import math
 import os
 import re
@@ -40,7 +39,7 @@
 from instructlab.dolomite.hf_models import GPTDolomiteForCausalLM
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-from transformers import AutoModelForCausalLM, get_scheduler
+from transformers import AutoConfig, AutoModelForCausalLM, get_scheduler
 import torch
 import torch.distributed
 
@@ -538,9 +537,8 @@ def main(args):
     tokenizer = setup_tokenizer(args.model_name_or_path, SPECIAL_TOKENS, CHAT_TEMPLATE)
     # device = torch.device("cuda", args.local_rank)
 
-    with open(Path(args.model_name_or_path) / "config.json") as conf_json:
-        model_conf = json.load(conf_json)
-    args.model_type = model_conf["model_type"]
+    model_conf = AutoConfig.from_pretrained(args.model_name_or_path)
+    args.model_type = model_conf.model_type
 
     #### distributed init #####
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))