TraceDetect-AI / download_text_data.py
nchdlhbctm's picture
Upload 13 files
646535a verified
Raw
History Blame Contribute Delete
1.8 kB
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 文件看看,里面的文本非常有意思!")