| import lmdb |
| import os |
|
|
| from tqdm import tqdm |
| from scripts.utils.others import TimeCounter |
|
|
|
|
| _10TB = 10995116277760 |
| _1TB = 1099511627776 |
| _1GB = 1073741824 |
|
|
|
|
| |
| def get_length(lmdb_dir): |
| env = lmdb.open(lmdb_dir, readonly=True, map_size=_1GB) |
| operator = env.begin() |
|
|
| length = int(operator.get('length'.encode()).decode()) |
|
|
| env.close() |
| return length |
|
|
|
|
| |
| def dump_lmdb(data_dict, lmdb_dir, verbose=True): |
| os.makedirs(lmdb_dir, exist_ok=True) |
|
|
| |
| env = lmdb.open(lmdb_dir, map_size=_10TB) |
|
|
| with env.begin(write=True) as operator: |
| if verbose: |
| iter_dict = tqdm(data_dict.items(), desc="Dumping data...") |
| else: |
| iter_dict = data_dict.items() |
|
|
| for k, v in iter_dict: |
| operator.put(key=str(k).encode(), value=str(v).encode()) |
|
|
| env.close() |
|
|
|
|
| |
| def jsonl2lmdb(jsonl_path, lmdb_dir): |
| os.makedirs(lmdb_dir, exist_ok=True) |
|
|
| |
| env = lmdb.open(lmdb_dir, map_size=_10TB) |
|
|
| with env.begin(write=True) as operator: |
| with TimeCounter("Loading data..."): |
| with open(jsonl_path, 'r', encoding='utf-8') as r: |
| cnt = 0 |
| for line in tqdm(r, desc="Parsing jsonl..."): |
| operator.put(key=str(cnt).encode(), value=line.encode()) |
| cnt += 1 |
|
|
| info = "Keys are as follows:\n" \ |
| " info: decription of dataset\n" \ |
| " length: length of data\n" \ |
| " 0 ~ length-1: index of each data\n" |
|
|
| operator.put(key='info'.encode(), value=info.encode()) |
| operator.put(key='length'.encode(), value=str(cnt).encode()) |
|
|
| env.close() |
