import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' # 设置镜像站点地址，确保从镜像站点下载数据集
# 注意，这个必须写在from transformers import BertTokenizer的前面，不然配置实现，有可能下载失败

from datasets import load_dataset
from transformers import BertTokenizer


# 设置数据集保存路径
dataset_path = "/root/autodl-tmp/dataset/Amazon_Reviews_CSV"

# 如果路径不存在，创建路径
os.makedirs(dataset_path, exist_ok=True)

# 下载并保存数据集
dataset = load_dataset("yassiracharki/Amazon_Reviews_for_Sentiment_Analysis_fine_grained_5_classes")


# 将每个分割（如 train, test）保存为 CSV 文件,to_csv将文件保存成CSV格式
# to_csv将文件保存成csv格式
for split in dataset.keys():
    dataset[split].to_csv(f"{dataset_path}/{split}.csv")

print(f"Dataset has been saved to {dataset_path}")