src/data_download.py

"""Download Amazon Reviews 2023 (Clothing_Shoes_and_Jewelry).

銆愰噸瑕佸彉鏇淬€戝畬鍏ㄦ姏寮?HuggingFace 鏁版嵁闆?API,鏀逛负鐩存帴浠?McAuley Lab 瀹樻柟
HTTPS 鏈嶅姟鍣ㄦ祦寮忔媺鍙栧師濮?jsonl.gz 鏂囦欢銆傚師鍥?
1. 鏂扮増 datasets 宸插簾寮?loading script, load_dataset(..., trust_remote_code=True)
   鎶?"Dataset scripts are no longer supported"銆?2. McAuley-Lab/Amazon-Reviews-2023 浠撳簱鐨?main 鍒嗘敮鍙湁鑴氭湰鍜?raw 鐩綍,
   parquet 鏂囦欢鍒嗘暎鍦?PR 鍒嗘敮(refs/pr/10)涓?涓?Clothing 绫荤洰骞朵笉瀹屾暣,
   渚濊禆 HF 璧板緱闈炲父鏇叉姌涓斿鏄撹俯鍧戙€?3. 瀹樻柟鏂囨。 (amazon-reviews-2023.github.io) 缁欑殑鐩撮摼鏈€绋?
   https://mcauleylab.ucsd.edu/public_datasets/data/amazon_2023/...
   涓€浠?jsonl.gz,娴佸紡 GET + gunzip + readline 涓変欢濂楀嵆鍙€?   瀵?sample 妯″紡灏ゅ叾鍙嬪ソ鈥斺€斾笅澶?N 鏉″氨鎺愭柇杩炴帴,涓嶇敤涓嬪畬鏁翠釜鏂囦欢銆?4. UCSD 鏈嶅姟鍣ㄥ湪鍥藉唴璁块棶閫氬父姣?huggingface.co 绋冲畾 (HF 缁忓父 SNI 闃绘柇)銆?
銆?026-05 淇銆?- pyarrow 鍐?parquet 鏃朵細鎸夊垪鎺ㄦ柇绫诲瀷, Amazon meta 鐨?price 瀛楁鏄?  "鍙兘鏁板瓧 / 鍙兘 '鈥? (em-dash) / 鍙兘 'None' / 鍙兘绌? 鐨勮剰娣峰悎瀛楁,
  pyarrow 璇曞浘鎺ㄦ垚 double 鏃堕亣鍒?'鈥? 鐩存帴 ArrowInvalid銆?  瑙ｅ喅: 鍐?parquet 鍓嶇敤 _clean_meta_for_parquet 鎶?price 寮鸿浆 float銆?- 鍚屾椂澧炲姞 jsonl 鍏滃簳鐩? meta 鎷夊畬鍏堝瓨涓€浠藉師濮?jsonl, 鍐嶅啓 parquet銆?  杩欐牱涓囦竴 parquet 鍙堝嚭鍒殑闂, 涓嬫鑳界绾ф仮澶? 涓嶅繀閲嶆柊鐖?25 鍒嗛挓銆?- reviews parquet 濡傛灉宸茬粡瀛樺湪, 璺宠繃 reviews 涓嬭浇闃舵, 鐩存帴澶嶇敤銆?"""
import gzip
import json
import logging
import os
from collections import defaultdict
from typing import Iterable, Iterator

import pandas as pd
import requests
from tqdm import tqdm

from . import config as cfg


logger = logging.getLogger(__name__)


# McAuley Lab 瀹樻柟鐩撮摼妯℃澘
_BASE = "https://mcauleylab.ucsd.edu/public_datasets/data/amazon_2023/raw"
_REVIEW_URL_TMPL = _BASE + "/review_categories/{cat}.jsonl.gz"
_META_URL_TMPL = _BASE + "/meta_categories/meta_{cat}.jsonl.gz"


def _stream_jsonl_gz(url: str) -> Iterator[dict]:
    """Stream a gzipped JSONL file and yield parsed rows.

    The function intentionally avoids downloading the full file to disk, so
    sample mode can stop early after enough valid records have been collected.
    """
    logger.info("GET %s", url)
    headers = {"User-Agent": "Mozilla/5.0 (acsa_clothing downloader)"}
    with requests.get(url, stream=True, headers=headers, timeout=(15, 60)) as r:
        r.raise_for_status()
        r.raw.decode_content = False
        with gzip.GzipFile(fileobj=r.raw) as gz:
            for raw_line in gz:
                if not raw_line:
                    continue
                try:
                    yield json.loads(raw_line)
                except json.JSONDecodeError as e:
                    logger.warning("Skipping one invalid JSON line: %s", e)
                    continue

def _records_to_df(records: Iterable[dict]) -> pd.DataFrame:
    return pd.DataFrame(list(records))


def _clean_meta_for_parquet(df: pd.DataFrame) -> pd.DataFrame:
    """鎶?meta 涓細璁?pyarrow 鎺ㄦ柇绫诲瀷澶辫触鐨勮剰鍊兼竻娲楁帀銆?
    Amazon 鍏冩暟鎹噷 price 鏄?"鍙兘鏁板瓧 / 鍙兘 '鈥? / 鍙兘 'None' / 鍙兘绌? 鐨勬贩鍚堝瓧娈?
    pandas 榛樿鎺ㄦ垚 object, pyarrow 钀?parquet 鏃跺嵈璇曞浘鎺ㄦ垚 double,
    閬囧埌 '鈥? (em-dash) 杩欑闈炴暟瀛楀瓧绗︿覆灏辨姏 ArrowInvalid銆?    鎵€浠ヨ繖閲屾墜鍔ㄦ妸 price 寮哄埗杞垚 float64, 鏃犳硶瑙ｆ瀽鐨勪竴寰嬪彉 NaN銆?    """
    if "price" in df.columns:
        df["price"] = pd.to_numeric(df["price"], errors="coerce")
    return df


def _backup_jsonl(records: list, parquet_path) -> str:
    """鍦ㄥ啓 parquet 鍓嶅厛钀戒竴浠藉師濮?jsonl 鍏滃簳銆?
    璺緞灏辨槸 parquet 鍚屽悕鎹㈠悗缂€銆傚嵆浣垮悗缁?to_parquet 澶辫触涔熶笉蹇呴噸鐖€?    """
    raw_jsonl_path = os.path.splitext(str(parquet_path))[0] + ".jsonl"
    os.makedirs(os.path.dirname(raw_jsonl_path) or ".", exist_ok=True)
    with open(raw_jsonl_path, "w", encoding="utf-8") as f:
        for row in records:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    logger.info("Backup jsonl saved to %s (%d rows)", raw_jsonl_path, len(records))
    return raw_jsonl_path


def download_sample(n_reviews: int = 30000, min_text_len: int = 20,
                    min_reviews_per_product: int = 2,
                    max_reviews_per_product: int = 5,
                    scan_multiplier: int = 5,
                    force: bool = False):
    """Sample 妯″紡: 娴佸紡璇诲彇鍓?N 鏉″彲鐢ㄨ瘎璁?鍐嶆媺瀵瑰簲 meta銆?
    'Usable' 鎸? text 闈炵┖涓旈暱搴?>= min_text_len, 涓?parent_asin 闈炵┖銆?    """
    review_url = _REVIEW_URL_TMPL.format(cat=cfg.CATEGORY)
    meta_url = _META_URL_TMPL.format(cat=cfg.CATEGORY)

    # ===== Reviews 闃舵 =====
    # 濡傛灉 reviews parquet 宸插瓨鍦?(涓婃璺戞垚鍔熶簡, 鏄?meta 闃舵宕╃殑),
    # 鐩存帴璇诲嚭鏉ュ鐢?seen_asins, 璺宠繃杩欐閲嶆柊涓嬭浇銆?
    if os.path.exists(cfg.RAW_REVIEWS_PATH) and not force:
        logger.info("Reviews parquet 宸插瓨鍦? 澶嶇敤: %s", cfg.RAW_REVIEWS_PATH)
        reviews_df = pd.read_parquet(cfg.RAW_REVIEWS_PATH)
        seen_asins = set(reviews_df["parent_asin"].dropna().unique().tolist())
        logger.info("浠庡凡鏈?reviews 涓仮澶?%d 鏉? %d 涓?unique parent_asins",
                    len(reviews_df), len(seen_asins))
    else:
        logger.info(
            "Streaming reviews (grouped sample, target=%d, min_per_product=%d, max_per_product=%d, scan_multiplier=%d) from %s ...",
            n_reviews, min_reviews_per_product, max_reviews_per_product, scan_multiplier, review_url,
        )
        candidate_target = max(n_reviews, n_reviews * max(scan_multiplier, 1))
        candidates = []
        pbar = tqdm(total=candidate_target, desc="review candidates")
        for row in _stream_jsonl_gz(review_url):
            text = (row.get("text") or "").strip()
            if len(text) < min_text_len:
                continue
            if not row.get("parent_asin"):
                continue
            candidates.append(row)
            pbar.update(1)
            if len(candidates) >= candidate_target:
                break
        pbar.close()

        by_asin = defaultdict(list)
        for row in candidates:
            by_asin[row["parent_asin"]].append(row)

        eligible = {asin: rows for asin, rows in by_asin.items() if len(rows) >= min_reviews_per_product}
        kept = []
        for asin, rows in sorted(eligible.items(), key=lambda x: len(x[1]), reverse=True):
            for row in rows[:max_reviews_per_product]:
                kept.append(row)
                if len(kept) >= n_reviews:
                    break
            if len(kept) >= n_reviews:
                break

        if len(kept) < n_reviews:
            used = {id(row) for row in kept}
            for row in candidates:
                if id(row) in used:
                    continue
                kept.append(row)
                if len(kept) >= n_reviews:
                    break

        reviews_df = _records_to_df(kept[:n_reviews])
        seen_asins = set(reviews_df["parent_asin"].dropna().unique().tolist())
        per_product = reviews_df.groupby("parent_asin").size()
        logger.info(
            "Got %d reviews, %d unique parent_asins, avg_reviews_per_product=%.2f, products_with_%d_plus_reviews=%d.",
            len(reviews_df), len(seen_asins), len(reviews_df) / max(len(seen_asins), 1),
            min_reviews_per_product, int((per_product >= min_reviews_per_product).sum()),
        )
        reviews_df.to_parquet(cfg.RAW_REVIEWS_PATH, index=False)

    # ===== Meta 闃舵 =====
    # 鍐嶆媺 meta,鍖归厤 parent_asin
    # 娉ㄦ剰: meta 娌℃湁淇濊瘉鎺掑簭,鐞嗚涓婅鎵畬鏁翠唤鎵嶈兘纭繚鍛戒腑鎵€鏈?asin銆?
    # 浣?meta 姣?reviews 灏忓緱澶?(Clothing meta 鍑犵櫨 MB 鍘嬬缉),瀹屾暣鎵竴閬嶅彲鎺ュ彈銆?    logger.info("Streaming meta to match parent_asins from %s ...", meta_url)
    meta_kept = []
    remaining = set(seen_asins)
    pbar = tqdm(total=len(seen_asins), desc="meta")
    for row in _stream_jsonl_gz(meta_url):
        asin = row.get("parent_asin")
        if asin in remaining:
            meta_kept.append(row)
            remaining.discard(asin)
            pbar.update(1)
            if not remaining:
                break  # 鍏ㄩ儴鍛戒腑,鎻愬墠鏂紑
    pbar.close()

    # 1) 鍏堣惤 jsonl 鍏滃簳, 鍝€曚笅闈?to_parquet 鐐镐簡涔熶笉蹇呴噸鐖?
    _backup_jsonl(meta_kept, cfg.RAW_META_PATH)

    # 2) 鍐嶅仛娓呮礂 + 鍐?parquet
    meta_df = _records_to_df(meta_kept)
    logger.info(
        "Got %d meta entries (matched %.1f%% of review asins).",
        len(meta_df),
        100 * len(meta_df) / max(len(seen_asins), 1),
    )
    meta_df = _clean_meta_for_parquet(meta_df)
    meta_df.to_parquet(cfg.RAW_META_PATH, index=False)
    return reviews_df, meta_df


def download_full():
    """Download the full review and metadata tables. Very large; use with care."""
    review_url = _REVIEW_URL_TMPL.format(cat=cfg.CATEGORY)
    meta_url = _META_URL_TMPL.format(cat=cfg.CATEGORY)

    logger.info("Downloading FULL reviews from %s ...", review_url)
    reviews = []
    for row in tqdm(_stream_jsonl_gz(review_url), desc="reviews"):
        reviews.append(row)
    reviews_df = _records_to_df(reviews)
    logger.info("Reviews: %d rows", len(reviews_df))
    reviews_df.to_parquet(cfg.RAW_REVIEWS_PATH, index=False)

    logger.info("Downloading FULL meta from %s ...", meta_url)
    metas = []
    for row in tqdm(_stream_jsonl_gz(meta_url), desc="meta"):
        metas.append(row)

    # 钀?jsonl 鍏滃簳
    _backup_jsonl(metas, cfg.RAW_META_PATH)

    meta_df = _records_to_df(metas)
    logger.info("Meta: %d rows", len(meta_df))
    meta_df = _clean_meta_for_parquet(meta_df)
    meta_df.to_parquet(cfg.RAW_META_PATH, index=False)
    return reviews_df, meta_df