from datasets import load_dataset
import pandas as pd
import os

print("正在连接 Hugging Face 拉取 HC3 中文语料库...")
# 使用 verification_mode="no_checks" 防止部分网络下元数据校验报错
dataset = load_dataset("Hello-SimpleAI/HC3-Chinese", name="all", split="train", verification_mode="no_checks", trust_remote_code=True)

human_texts = []
ai_texts = []
target_count = 500  # 咱们先各抓 500 条用来快速跑通流程

print(f"开始清洗数据，目标：提取 {target_count}条人类文本 + {target_count}条AI文本...")

for item in dataset:
    # 提取人类回答 (剔除太短的废话)
    if len(item['human_answers']) > 0 and len(item['human_answers'][0]) > 20:
        if len(human_texts) < target_count:
            human_texts.append(item['human_answers'][0])

    # 提取 ChatGPT 回答
    if len(item['chatgpt_answers']) > 0 and len(item['chatgpt_answers'][0]) > 20:
        if len(ai_texts) < target_count:
            ai_texts.append(item['chatgpt_answers'][0])

    if len(human_texts) >= target_count and len(ai_texts) >= target_count:
        break

# 制作结构化的 DataFrame (标签规则统一：0代表AI，1代表真实)
df_human = pd.DataFrame({'text': human_texts, 'label': 1})
df_ai = pd.DataFrame({'text': ai_texts, 'label': 0})

# 合并并打乱顺序 (Shuffle)
df_all = pd.concat([df_human, df_ai]).sample(frac=1, random_state=42).reset_index(drop=True)

# 确保 data 文件夹存在
os.makedirs("./data", exist_ok=True)
csv_path = "./data/text_dataset.csv"
df_all.to_csv(csv_path, index=False, encoding='utf-8')

print(f"✅ 语料库构建完毕！已保存为 {csv_path} (共 {len(df_all)} 条数据)")
print("你可以打开这个 csv 文件看看，里面的文本非常有意思！")