235 lines
9.8 KiB
Python
235 lines
9.8 KiB
Python
|
import argparse
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
import shutil
|
|||
|
|
import h5py
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
#HDFS_INDEX_DATA_PATH = "/user/tc_agi/llm/index_datasets/"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def calcu_data_length(data, language):
|
|||
|
|
if language == "zh":
|
|||
|
|
return 0.7 * len(data)
|
|||
|
|
if language in ["en", "code"]:
|
|||
|
|
return 1.5 * len(data.split(" "))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def update_data(length, return_res):
|
|||
|
|
if length < 4000:
|
|||
|
|
return_res["less_4k"] += 1
|
|||
|
|
if length >= 4000 and length < 8000:
|
|||
|
|
return_res["4k-8k"] += 1
|
|||
|
|
if length >= 8000 and length < 16000:
|
|||
|
|
return_res["8k-16k"] += 1
|
|||
|
|
if length >= 16000 and length < 32000:
|
|||
|
|
return_res["16k-32k"] += 1
|
|||
|
|
if length >= 32000 and length < 64000:
|
|||
|
|
return_res["32k-64k"] += 1
|
|||
|
|
if length >= 64000 and length < 128000:
|
|||
|
|
return_res["64k-128k"] += 1
|
|||
|
|
if length >= 128000 and length < 256000:
|
|||
|
|
return_res["128k-256k"] += 1
|
|||
|
|
if length >= 256000:
|
|||
|
|
return_res["more_256k"] += 1
|
|||
|
|
return return_res
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 限制sample样本的大小
|
|||
|
|
def check_string_length(json_obj, list_limit_length=1, str_limit_length=50):
|
|||
|
|
if isinstance(json_obj, dict):
|
|||
|
|
for key, value in json_obj.items():
|
|||
|
|
if isinstance(value, list) and len(value) > list_limit_length:
|
|||
|
|
json_obj[key] = value[:list_limit_length] # 截断列表长度
|
|||
|
|
elif isinstance(value, str):
|
|||
|
|
try:
|
|||
|
|
value_json = json.loads(value) # 尝试解析字符串为 JSON 对象
|
|||
|
|
check_string_length(value_json, list_limit_length, str_limit_length) # 递归处理 JSON 对象
|
|||
|
|
json_obj[key] = json.dumps(value_json,ensure_ascii=False) # 将修改后的 JSON 对象转回字符串
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
if len(value) > str_limit_length:
|
|||
|
|
json_obj[key] = value[:str_limit_length] + "..." # 截断字符串长度
|
|||
|
|
elif isinstance(value, dict):
|
|||
|
|
check_string_length(value, list_limit_length, str_limit_length) # 递归遍历嵌套的字典或列表
|
|||
|
|
elif isinstance(json_obj, list):
|
|||
|
|
for i in range(len(json_obj)):
|
|||
|
|
if isinstance(json_obj[i], list) and len(json_obj[i]) > list_limit_length:
|
|||
|
|
json_obj[i] = json_obj[i][:list_limit_length] # 截断列表长度
|
|||
|
|
elif isinstance(json_obj[i], str):
|
|||
|
|
try:
|
|||
|
|
value_json = json.loads(json_obj[i]) # 尝试解析字符串为 JSON 对象
|
|||
|
|
check_string_length(value_json, list_limit_length, str_limit_length) # 递归处理 JSON 对象
|
|||
|
|
json_obj[i] = json.dumps(value_json,ensure_ascii=False) # 将修改后的 JSON 对象转回字符串
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
if len(json_obj[i]) > str_limit_length:
|
|||
|
|
json_obj[i] = json_obj[i][:str_limit_length] + "..." # 截断字符串长度
|
|||
|
|
elif isinstance(json_obj[i], dict):
|
|||
|
|
check_string_length(json_obj[i], list_limit_length, str_limit_length) # 递归遍历嵌套的字典或列表
|
|||
|
|
|
|||
|
|
|
|||
|
|
def check_index_line(data_path:str, index_path:str, h5_path):
|
|||
|
|
index_len = 0
|
|||
|
|
data_len = 0
|
|||
|
|
if os.path.exists(data_path):
|
|||
|
|
with open(data_path, 'r') as file:
|
|||
|
|
# 逐行迭代文件对象
|
|||
|
|
data_len = sum(1 for line in file)
|
|||
|
|
if os.path.exists(index_path):
|
|||
|
|
with open(data_path, 'r') as file:
|
|||
|
|
# 逐行迭代文件对象
|
|||
|
|
index_len = sum(1 for line in file)
|
|||
|
|
|
|||
|
|
if index_len == data_len and index_len != 0:
|
|||
|
|
return index_len, data_len, True
|
|||
|
|
else:
|
|||
|
|
return index_len, data_len, False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def list_files_in_directory(directory):
|
|||
|
|
file_list = []
|
|||
|
|
|
|||
|
|
# 遍历指定目录下的所有文件
|
|||
|
|
for root, dirs, files in os.walk(directory):
|
|||
|
|
for file in files:
|
|||
|
|
if file.endswith(('.json', '.jsonl')) and "meta.json" not in file:
|
|||
|
|
file_path = os.path.join(root, file)
|
|||
|
|
file_list.append(file_path)
|
|||
|
|
|
|||
|
|
return file_list
|
|||
|
|
|
|||
|
|
def build_index(file_list, output_path, language):
|
|||
|
|
offset = 0
|
|||
|
|
if not os.path.exists(output_path):
|
|||
|
|
os.mkdir(output_path)
|
|||
|
|
meta_path = os.path.join(output_path, "meta.json")
|
|||
|
|
data_path = os.path.join(output_path, "data.jsonl")
|
|||
|
|
nlines = 0
|
|||
|
|
data_sample = None
|
|||
|
|
all_token = 0
|
|||
|
|
starts = []
|
|||
|
|
length_distribute = {
|
|||
|
|
"less_4k": 0,
|
|||
|
|
"4k-8k": 0,
|
|||
|
|
"8k-16k": 0,
|
|||
|
|
"16k-32k": 0,
|
|||
|
|
"32k-64k": 0,
|
|||
|
|
"64k-128k": 0,
|
|||
|
|
"128k-256k": 0,
|
|||
|
|
"more_256k": 0,
|
|||
|
|
}
|
|||
|
|
w_index_h5 = h5py.File(os.path.join(output_path, "index.h5"), "w")
|
|||
|
|
w_index = open(os.path.join(output_path, "index"), "w")
|
|||
|
|
|
|||
|
|
print(data_path)
|
|||
|
|
if os.path.exists(data_path):
|
|||
|
|
print("data.jsonl already exists! No need create data.jsonl!")
|
|||
|
|
with open(data_path, "rb") as fin:
|
|||
|
|
for idx, line in tqdm(enumerate(fin)):
|
|||
|
|
nlines += 1
|
|||
|
|
offset += len(line)
|
|||
|
|
starts.append(offset)
|
|||
|
|
w_index.write(f"{offset}\n")
|
|||
|
|
decoded_line = line.decode("utf-8")
|
|||
|
|
if idx == 0:
|
|||
|
|
data_sample = json.loads(decoded_line)
|
|||
|
|
check_string_length(data_sample)
|
|||
|
|
length = calcu_data_length(decoded_line, language)
|
|||
|
|
all_token += length
|
|||
|
|
update_dic = update_data(length, length_distribute)
|
|||
|
|
length_distribute.update(update_dic)
|
|||
|
|
else:
|
|||
|
|
print("data.jsonl not exists! Creating data.jsonl...")
|
|||
|
|
w_data_path = open(data_path, "wb")
|
|||
|
|
|
|||
|
|
#for ds in tqdm(file_list):
|
|||
|
|
if 1:
|
|||
|
|
with open(file_list, "rb") as fin:
|
|||
|
|
for idx, line in tqdm(enumerate(fin)):
|
|||
|
|
nlines += 1
|
|||
|
|
offset += len(line)
|
|||
|
|
starts.append(offset)
|
|||
|
|
w_index.write(f"{offset}\n")
|
|||
|
|
decoded_line = line.decode("utf-8")
|
|||
|
|
if idx == 0:
|
|||
|
|
data_sample = json.loads(decoded_line)
|
|||
|
|
check_string_length(data_sample)
|
|||
|
|
length = calcu_data_length(decoded_line, language)
|
|||
|
|
all_token += length
|
|||
|
|
update_dic = update_data(length, length_distribute)
|
|||
|
|
length_distribute.update(update_dic)
|
|||
|
|
w_data_path.write(line)
|
|||
|
|
w_data_path.close()
|
|||
|
|
|
|||
|
|
if nlines == 0:
|
|||
|
|
return 0, 0, False
|
|||
|
|
|
|||
|
|
assert os.path.exists(output_path), f"Jsonline dataset '{data_path}' not found."
|
|||
|
|
meta_dic = {
|
|||
|
|
"language": language,
|
|||
|
|
"nlines": nlines,
|
|||
|
|
"nbytes": offset,
|
|||
|
|
"length_distribute": length_distribute,
|
|||
|
|
"avg_token_per_line": all_token / nlines,
|
|||
|
|
"data_sample": data_sample,
|
|||
|
|
}
|
|||
|
|
with open(meta_path, "w") as w:
|
|||
|
|
w.write(json.dumps(meta_dic, ensure_ascii=False) + "\n")
|
|||
|
|
|
|||
|
|
w_index_h5.create_dataset("index", data=starts)
|
|||
|
|
|
|||
|
|
w_index.close()
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# python trans.py -p /home/jeeves/souyun/ -l zh -o /home/jeeves/souyunout -hdfs souyunout
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--path", "-p", required=True, help="Input data path.")
|
|||
|
|
parser.add_argument("--language", "-l", required=True, help="language type, select 'zh' or 'en'")
|
|||
|
|
parser.add_argument("--output", "-o", required=True, help="Output index Data dir path.")
|
|||
|
|
#parser.add_argument("--hdfs_name", "-hdfs", required=True, help="hdfs save name. hdfs dir is '/user/tc_agi/llm/index_datasets/'")
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
#file_list = list_files_in_directory(args.path)
|
|||
|
|
|
|||
|
|
file_list = args.path
|
|||
|
|
#print(file_list)
|
|||
|
|
# if os.path.isfile(args.path):
|
|||
|
|
# file_list = [args.path]
|
|||
|
|
# else:
|
|||
|
|
# for file_namei in os.listdir(args.path):
|
|||
|
|
# tmp_dir = os.path.join(args.path, file_namei)
|
|||
|
|
# if os.path.isfile(tmp_dir):
|
|||
|
|
# file_list.append(tmp_dir)
|
|||
|
|
# else:
|
|||
|
|
# for file_name_2 in os.listdir(tmp_dir):
|
|||
|
|
# file_list.append(os.path.join(tmp_dir, file_name_2))
|
|||
|
|
#file_list.sort()
|
|||
|
|
|
|||
|
|
|
|||
|
|
if os.path.exists(f"{args.output}"):
|
|||
|
|
shutil.rmtree(f"{args.output}")
|
|||
|
|
|
|||
|
|
build_index(file_list, args.output, args.language)
|
|||
|
|
|
|||
|
|
output_jsonl = f"{args.output}/data.jsonl"
|
|||
|
|
output_index = f"{args.output}/index"
|
|||
|
|
output_h5 = f"{args.output}/index.h5"
|
|||
|
|
output_meta = f"{args.output}/meta.json"
|
|||
|
|
index_len, data_len, check_result = check_index_line(output_jsonl, output_index, output_h5)
|
|||
|
|
'''
|
|||
|
|
if check_result:
|
|||
|
|
#上传hdfs
|
|||
|
|
os.system("hadoop fs -get -f hdfs://hdfs01:8020/etc/infra-conf-admin/packages/hdfs-copy-1.0-SNAPSHOT.jar .")
|
|||
|
|
os.system("export HADOOP_USER_NAME=tc_agi")
|
|||
|
|
os.system("export HADOOP_USER_PASSWORD=IH2U3AS1D")
|
|||
|
|
os.system(f"hadoop jar hdfs-copy-1.0-SNAPSHOT.jar com.zhihu.platform.Main --dst hdfs://hdfs01:8020{HDFS_INDEX_DATA_PATH}{args.hdfs_name}/data.jsonl --src file://{output_jsonl} --parallelism 50 --dstToken tc_agi-IH2U3AS1D")
|
|||
|
|
os.system(f"hadoop jar hdfs-copy-1.0-SNAPSHOT.jar com.zhihu.platform.Main --dst hdfs://hdfs01:8020{HDFS_INDEX_DATA_PATH}{args.hdfs_name}/index --src file://{output_index} --parallelism 50 --dstToken tc_agi-IH2U3AS1D")
|
|||
|
|
os.system(f"hadoop jar hdfs-copy-1.0-SNAPSHOT.jar com.zhihu.platform.Main --dst hdfs://hdfs01:8020{HDFS_INDEX_DATA_PATH}{args.hdfs_name}/index.h5 --src file://{output_h5} --parallelism 50 --dstToken tc_agi-IH2U3AS1D")
|
|||
|
|
os.system(f"hadoop jar hdfs-copy-1.0-SNAPSHOT.jar com.zhihu.platform.Main --dst hdfs://hdfs01:8020{HDFS_INDEX_DATA_PATH}{args.hdfs_name}/meta.json --src file://{output_meta} --parallelism 50 --dstToken tc_agi-IH2U3AS1D")
|
|||
|
|
print(f"index文件行数：{index_len}行，数据文件行数：{data_len}行")
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
logging.error("==========data.jsonl和data_index行数不一致==========")
|
|||
|
|
logging.error(f"index文件行数：{index_len}行，数据文件行数：{data_len}行")
|
|||
|
|
'''
|