from datasets import Dataset, concatenate_datasets

data_path = "../../LSPD/train/spectrum"

import os
from os import path

dataset_list = []

for folder in os.listdir(data_path):
    folder_path = path.join(data_path, folder)
    print(f"正在加载文件夹: {folder_path}")

    dataset_list.append(Dataset.load_from_disk(folder_path))
    print(f"已加载 {folder}, data length: {len(dataset_list[-1])}")

    if len(dataset_list) >= 5:  # 只加载前2个文件夹以节省内存
        break

# 合并数据集
print("正在合并数据集...")
combined_dataset = concatenate_datasets(dataset_list)
print(f"合并后的数据集长度: {len(combined_dataset)}")
