加载数据集,支持 JSON、JSONL 和 Parquet 文件格式,并始终返回 list。 参数: dataset_path (str): 数据集文件路径。 dataset (list, optional): 如果提供了 dataset,则直接返回。 返回: list: 加载的数据集,统一为列表格式。
(dataset_path, dataset=None)
| 33 | |
| 34 | |
| 35 | def 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 | |
| 88 | class BaseEvaluator: |
| 89 | def __init__( |
no test coverage detected