github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

walk folder to get dataset_type_ids_map

pull/477/head
gaoyang07 2023-11-07 17:13:45 +08:00
parent 4d1b1cd5f1
commit 61f953bb7b
2 changed files with 15 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
internlm/data/packed_dataset.py
View File

@ -15,7 +15,7 @@ from tqdm import tqdm
from internlm.core.context import global_context as gpc from internlm.core.context import global_context as gpc
from internlm.data.single_dataset import JsonlDataset from internlm.data.single_dataset import JsonlDataset
from internlm.data.utils import get_dataset_type_id from internlm.data.utils import get_dataset_type_id, get_dataset_type_ids_map
from internlm.utils.logger import get_logger from internlm.utils.logger import get_logger
DEFAULT_SEED = 1024 DEFAULT_SEED = 1024
@ -373,6 +373,8 @@ def get_packed_dataset_without_short_length(
datasets = [] datasets = []
delete_samples = 0 delete_samples = 0
DATASET_TYPE_IDS_MAP = get_dataset_type_ids_map(folder)
if gpc.get_global_rank() == 0: if gpc.get_global_rank() == 0:
triples = [list(os.walk(folder, followlinks=True))] triples = [list(os.walk(folder, followlinks=True))]
else: else:
@ -400,7 +402,7 @@ def get_packed_dataset_without_short_length(
len(catch_ml_keys) < 2 len(catch_ml_keys) < 2
), f"The file name `{fp}` matched the following resample keys:{catch_ml_keys}" ), f"The file name `{fp}` matched the following resample keys:{catch_ml_keys}"
ds_type_id = get_dataset_type_id(path=fp) ds_type_id = get_dataset_type_id(DATASET_TYPE_IDS_MAP, path=fp)
ds = JsonlDataset(fp, ds_type_id, min_length=min_length_num) ds = JsonlDataset(fp, ds_type_id, min_length=min_length_num)
if hasattr(ds, "old_length"): if hasattr(ds, "old_length"):

17
internlm/data/utils.py
View File

@ -1,21 +1,26 @@
#!/usr/bin/env python #!/usr/bin/env python
# -*- encoding: utf-8 -*- # -*- encoding: utf-8 -*-
import os
import re
import torch import torch
from internlm.core.context import global_context as gpc from internlm.core.context import global_context as gpc
DATASET_TYPE_IDS_MAP = {"en": 0, "cn": 1, "code": 2}
def get_dataset_type_ids_map(path):
dirlist = list(os.listdir(path))
dirlist.sort()
return {key: idx for idx, key in enumerate(dirlist)}
def get_dataset_type_id(path): def get_dataset_type_id(dataset_type_ids_map, path):
import re
match_idxes = [] match_idxes = []
for key, idx in DATASET_TYPE_IDS_MAP.items():
for key, idx in dataset_type_ids_map.items():
if re.search(rf"/[z_]*{key}/", path): if re.search(rf"/[z_]*{key}/", path):
match_idxes.append(idx) match_idxes.append(idx)
assert len(match_idxes) == 1, f"{path}, match_idxes should be 1, but got {match_idxes} from {DATASET_TYPE_IDS_MAP}" assert len(match_idxes) == 1, f"{path}, match_idxes should be 1, but got {match_idxes} from {dataset_type_ids_map}"
return match_idxes[0] return match_idxes[0]