# 依赖：pip install datasets pandas

import pandas as pd
from datasets import load_dataset

def download_sst2_to_csv():
    """下载SST-2数据集并保存为CSV文件"""
    print("📥 正在下载SST-2数据集...")
    
    try:
        # 直接从huggingface下载SST-2
        dataset = load_dataset("sst2")
    except:
        # 备用方案：从GLUE下载
        dataset = load_dataset("glue", "sst2")
    
    print("✅ 数据集下载成功！")
    
    # 转换并保存每个分割
    for split_name, split_data in dataset.items():
        # 转换为DataFrame
        df_data = []
        for item in split_data:
            df_data.append({
                "text": item["sentence"],
                "label": item["label"]
            })
        
        df = pd.DataFrame(df_data)
        
        # 保存为CSV
        filename = f"sst2_{split_name}.csv"
        df.to_csv(filename, index=False, encoding="utf-8")
        
        print(f"💾 已保存: {filename} ({len(df)} 条样本)")
    
    print("🎉 完成！生成文件:")
    print("  - sst2_train.csv")
    print("  - sst2_validation.csv") 
    print("  - sst2_test.csv")

if __name__ == "__main__":
    download_sst2_to_csv()