load_data.py

"""Step 1: load a REAL Amazon product catalog -> docs.json.

Tries several public datasets, auto-detects title/description columns,
and streams just N rows (no giant download). Fully variety-rich data.
"""
import json
from datasets import load_dataset

N = 10000
CANDIDATES = [
    "ckandemir/amazon-products",
    "cvnberk/amazon-products",
    "bprateek/amazon_product_description",
]
TEXT_HINTS = ["title", "product", "name", "description", "about", "feature", "text"]


def pick_fields(example):
    keys = list(example.keys())

    def find(hints):
        for k in keys:
            if any(h in k.lower() for h in hints):
                return k
        return None

    title = find(["title", "product name", "name"])
    desc = find(["description", "about", "feature"])
    chosen = [c for c in [title, desc] if c]
    if not chosen:
        chosen = [k for k in keys if any(h in k.lower() for h in TEXT_HINTS)]
    return chosen


ds = None
for name in CANDIDATES:
    try:
        print("Trying", name, "...")
        ds = load_dataset(name, split="train", streaming=True)
        print("Loaded", name)
        break
    except Exception as e:
        print("  failed:", e)

if ds is None:
    raise SystemExit("Could not load any dataset. Check internet / datasets version.")

docs, seen = [], set()
fields = None
for row in ds:
    if fields is None:
        fields = pick_fields(row)
        print("Using fields:", fields)
    parts = []
    for f in fields:
        v = row.get(f)
        if isinstance(v, str) and v.strip():
            parts.append(v.strip())
        elif isinstance(v, list):
            parts.append(" ".join(str(x) for x in v if x))
    text = " ".join((" — ".join(parts)).split())
    if len(text) < 20 or text in seen:
        continue
    seen.add(text)
    docs.append(text[:500])
    if len(docs) >= N:
        break

json.dump(docs, open("docs.json", "w"))
print(f"Saved {len(docs)} product docs -> docs.json")
print("Sample:", docs[0][:160])