MCPcopy Create free account
hub / github.com/InternLM/InternBootcamp / load_dataset

Function load_dataset

internbootcamp/src/base_evaluator.py:35–86  ·  view source on GitHub ↗

加载数据集,支持 JSON、JSONL 和 Parquet 文件格式,并始终返回 list。 参数: dataset_path (str): 数据集文件路径。 dataset (list, optional): 如果提供了 dataset,则直接返回。 返回: list: 加载的数据集,统一为列表格式。

(dataset_path, dataset=None)

Source from the content-addressed store, hash-verified

33
34
35def load_dataset(dataset_path, dataset=None):
36 """
37 加载数据集,支持 JSON、JSONL 和 Parquet 文件格式,并始终返回 list。
38
39 参数:
40 dataset_path (str): 数据集文件路径。
41 dataset (list, optional): 如果提供了 dataset,则直接返回。
42
43 返回:
44 list: 加载的数据集,统一为列表格式。
45 """
46 if dataset_path and not dataset:
47 # 获取文件扩展名
48 _, ext = os.path.splitext(dataset_path)
49 ext = ext.lower() # 统一转换为小写
50
51 if ext == ".json":
52 # 加载 JSON 文件
53 with open(dataset_path, "r", encoding="utf-8") as f:
54 data = json.load(f)
55 # 确保返回的是列表
56 dataset = data if isinstance(data, list) else [data]
57
58 elif ext == ".jsonl":
59 # 加载 JSONL 文件
60 dataset = []
61 with jsonlines.open(dataset_path) as reader:
62 for line in reader:
63 dataset.append(line)
64
65 elif ext == ".parquet":
66 # 加载 Parquet 文件并转换为列表
67 df = pd.read_parquet(dataset_path)
68 dataset = df.to_dict(orient='records') # 转换为字典列表
69
70 # 处理 parquet 加载后的数据类型问题
71 for item in dataset:
72 # 确保 messages 和 prompt 字段是 Python 列表而不是 numpy 数组
73 if 'messages' in item and hasattr(item['messages'], 'tolist'):
74 item['messages'] = item['messages'].tolist()
75 elif 'messages' in item and not isinstance(item['messages'], list):
76 item['messages'] = list(item['messages'])
77
78 if 'prompt' in item and hasattr(item['prompt'], 'tolist'):
79 item['prompt'] = item['prompt'].tolist()
80 elif 'prompt' in item and not isinstance(item['prompt'], list):
81 item['prompt'] = list(item['prompt'])
82
83 else:
84 raise ValueError(f"不支持的文件格式: {ext}")
85
86 return dataset
87
88class BaseEvaluator:
89 def __init__(

Callers 6

create_evaluation_taskFunction · 0.90
run_evaluationMethod · 0.90
run_evaluationMethod · 0.85
generate_sft_datasetFunction · 0.85
generate_rm_datasetFunction · 0.85
generate_rl_datasetFunction · 0.85

Calls 3

appendMethod · 0.80
loadMethod · 0.45
to_dictMethod · 0.45

Tested by

no test coverage detected