Direction9/database/test_download.py

#!/usr/bin/env python
"""
快速测试脚本 - 下载单个领域的小样本验证链路

用法:
    python test_download.py
"""

import os
import re
import json
import random
from pathlib import Path

# 设置 HuggingFace 镜像
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from datasets import load_dataset


def preprocess_text(text: str, lang: str = "en") -> str | None:
    """预处理文本"""
    text = text.strip()
    if not text:
        return None

    if lang == "en":
        text = re.sub(r'[^\w\s.,!?;:()"\']', '', text)
        words = text.split()
        if len(words) < 50 or len(words) > 200:
            return None
        text = ' '.join(words[:200]).strip()
    else:
        text = re.sub(r'[^\u4e00-\u9fff0-9.,!?;:()""\']', '', text)
        if len(text) < 100 or len(text) > 500:
            return None

    return text if text else None


def main():
    print("=" * 50)
    print("快速测试：下载 CNN/DailyMail 英文新闻")
    print("=" * 50)

    # 加载数据集
    print("\n加载数据集...")
    ds = load_dataset("cnn_dailymail", "3.0.0", split="train")

    # 随机采样 20 条
    print("采样并预处理...")
    texts = [item["article"] for item in random.sample(list(ds), 20)]

    # 预处理
    corpus = []
    for text in texts:
        processed = preprocess_text(text, lang="en")
        if processed and processed not in corpus:
            corpus.append(processed)
        if len(corpus) >= 10:
            break

    print(f"成功提取 {len(corpus)} 条有效文本")

    # 保存
    output_path = Path("database/corpus/news_en_test/texts.jsonl")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w", encoding="utf-8") as f:
        for i, text in enumerate(corpus):
            record = {
                "text": text,
                "metadata": {"index": i, "domain": "news_en_test"}
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    print(f"已保存至：{output_path}")

    # 验证读取
    print("\n验证读取...")
    with open(output_path, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            record = json.loads(line)
            text = record["text"]
            print(f"  [{i}] 长度：{len(text)} 字符")
            if i >= 2:
                break

    print("\n" + "=" * 50)
    print("测试完成！链路验证通过")
    print("=" * 50)


if __name__ == "__main__":
    main()