File size: 1,776 Bytes
c6e7238 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import json
from pathlib import Path
from collections import defaultdict
DATASETS = {}
for path in Path("configs/dataset_configs").glob("*.json"):
dataset_id = path.stem
DATASETS[dataset_id] = json.loads(path.read_text())
def fetch_model_params(model):
model_path = model if model.endswith(".json") else f"configs/{model}.json"
with open(model_path) as f:
params = json.load(f)
dataset_ids = []
for d in params.get("datasets"):
if isinstance(d, list):
dataset_ids.append(d[0])
else:
dataset_ids.append(d)
no_datasets = params.get("no_dataset", False)
assert no_datasets or len(dataset_ids) > 0, "You must specify at least one dataset id in the model config"
datasets = {}
last_dataset = None
for dataset_id in dataset_ids:
assert dataset_id in DATASETS, f"Dataset '{dataset_id}' was not found under dataset_configs/ folder. Please follow the example.json in that folder."
dataset = DATASETS[dataset_id]
assert params["n_vocab"] >= dataset["n_vocab"], f"The embedding table size '{params['n_vocab']}' must be greater or equal to the vocab size used to encode the dataset '{dataset_id}' ({dataset['n_vocab']})"
datasets[dataset_id] = dataset
last_dataset = dataset
if last_dataset is not None:
params["padding_id"] = last_dataset.get("padding_id", 0)
params["eos_id"] = last_dataset.get("eos_id", 1)
params["dataset_configs"] = datasets
# Set some other parameter defaults
params["mlm_training"] = params.get("mlm_training") == True
params["causal"] = not params["mlm_training"]
# Set all other parameter values to default to None
params = defaultdict(lambda: None, params)
return params
|