Upload brain_virality_predictor/downloader.py with huggingface_hub

brain_virality_predictor/downloader.py

@@ -0,0 +1,191 @@
+"""
+Video download orchestrator.
+Consumes Apify-style JSON dumps (good/okish/bad) + cookies.txt,
+uses yt-dlp to fetch videos into labeled folders.
+"""
+
+import json, random, re, subprocess, time
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+
+def parse_json_labels(json_dir: Path) -> Dict[str, Path]:
+    """
+    Scan a directory for JSON label files.
+    Accepts patterns like good*.json, bad*.json, okish*.json, neutral*.json.
+    """
+    candidates = {lbl: [] for lbl in ("good", "okish", "bad")}
+    for p in json_dir.glob("*.json"):
+        name = p.name.lower()
+        if re.search(r"\(\d+\)", name):  # skip duplicates like "(1)"
+            continue
+        if "good" in name or "banger" in name or "viral" in name or "winner" in name:
+            candidates["good"].append(p)
+        elif "bad" in name or "flop" in name or "loser" in name or "poor" in name:
+            candidates["bad"].append(p)
+        elif "okish" in name or "neutral" in name or "avg" in name or "mid" in name or "average" in name:
+            candidates["okish"].append(p)
+
+    chosen = {}
+    for lbl, paths in candidates.items():
+        if paths:
+            chosen[lbl] = min(paths, key=lambda p: (len(p.name), p.name))
+    return chosen
+
+
+def build_job_queue(json_paths: Dict[str, Path], out_root: Path) -> Tuple[List, List]:
+    """
+    Build deduplicated download jobs from JSON files.
+    Returns: (ig_jobs, other_jobs) where each job is (url, target_path, label).
+    """
+    out_root = Path(out_root)
+    out_root.mkdir(parents=True, exist_ok=True)
+
+    def safe_name(s, n=30):
+        return "".join(c if c.isalnum() or c in "-_" else "_" for c in str(s))[:n]
+
+    def out_path(label, item, idx):
+        handle = safe_name(item.get("handle") or "unknown", 30)
+        score = f"{item.get('outlier_score', 0):.2f}".replace(".", "p")
+        return out_root / label / f"{idx:03d}_{score}_{handle}.mp4"
+
+    seen_urls = set()
+    ig_jobs, other_jobs = [], []
+
+    for label, json_path in json_paths.items():
+        items = json.load(open(json_path))
+        for idx, item in enumerate(items):
+            url = (item.get("url") or "").strip()
+            if not url or url in seen_urls:
+                continue
+            seen_urls.add(url)
+            target = out_path(label, item, idx)
+            platform = (item.get("platform") or "").lower()
+            is_ig = "instagram" in platform or "instagram.com" in url
+            (ig_jobs if is_ig else other_jobs).append((url, target, label))
+
+    return ig_jobs, other_jobs
+
+
+def categorize_error(err: str) -> str:
+    if not err:
+        return "unknown"
+    t = err.lower()
+    if "rate" in t or "429" in t:
+        return "rate_limited"
+    if "login" in t or "private" in t:
+        return "private"
+    if "not found" in t or "404" in t or "removed" in t or "unavailable" in t:
+        return "deleted"
+    if "timeout" in t:
+        return "timeout"
+    return "other"
+
+
+def download_one(url: str, target: Path, cookies: Path, timeout: int = 240) -> Dict:
+    if target.exists() and target.stat().st_size > 100_000:
+        return {"url": url, "path": str(target), "status": "skip", "error": None, "category": None}
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    cmd = [
+        "yt-dlp", url,
+        "-f", "best[ext=mp4]/mp4/best",
+        "-o", str(target),
+        "--no-warnings", "--quiet",
+        "--retries", "5",
+        "--socket-timeout", "30",
+        "--no-playlist",
+        "--max-filesize", "200M",
+    ]
+    if cookies.exists():
+        cmd += ["--cookies", str(cookies)]
+
+    try:
+        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+        if r.returncode == 0 and target.exists() and target.stat().st_size > 100_000:
+            return {"url": url, "path": str(target), "status": "ok", "error": None, "category": None}
+        err = (r.stderr or r.stdout or "?")[-300:]
+        return {"url": url, "path": str(target), "status": "fail",
+                "error": err, "category": categorize_error(err)}
+    except subprocess.TimeoutExpired:
+        return {"url": url, "path": str(target), "status": "fail",
+                "error": "timeout", "category": "timeout"}
+    except Exception as e:
+        return {"url": url, "path": str(target), "status": "fail",
+                "error": str(e), "category": categorize_error(str(e))}
+
+
+def download_all(json_dir: Path, cookies_path: Path, out_root: Path,
+                 workers: int = 4, log_path: Path = None) -> List[Dict]:
+    """
+    Full orchestrated download pipeline.
+    Returns list of result dicts.
+    """
+    json_paths = parse_json_labels(json_dir)
+    assert len(json_paths) == 3, f"Need 3 label JSONs (good/okish/bad), got: {list(json_paths.keys())}"
+
+    ig_jobs, other_jobs = build_job_queue(json_paths, out_root)
+    total = len(ig_jobs) + len(other_jobs)
+    print(f"Queue: {len(other_jobs)} TT/YT  +  {len(ig_jobs)} IG  =  {total} unique URLs")
+
+    cookies = Path(cookies_path)
+    has_cookies = cookies.exists() and "sessionid" in cookies.read_text(errors="ignore").lower()
+    print(f"Cookies: {'logged in' if has_cookies else 'missing — IG may fail'}")
+
+    results: List[Dict] = []
+
+    # Phase 1: TT/YT in parallel
+    print(f"\nPhase 1: TikTok + YouTube Shorts (parallel ×{workers})")
+    t0 = time.perf_counter()
+    with ThreadPoolExecutor(max_workers=workers) as ex:
+        futs = {ex.submit(download_one, u, t, cookies): (u, t, l) for u, t, l in other_jobs}
+        for i, fut in enumerate(as_completed(futs), 1):
+            res = fut.result()
+            results.append(res)
+            icon = {"ok": "✓", "skip": "→", "fail": "✗"}.get(res["status"], "?")
+            if i % 10 == 0 or res["status"] == "fail":
+                print(f"  [{i}/{len(other_jobs)}] {icon} {Path(res['path']).name[:40]}")
+    print(f"Phase 1 done in {(time.perf_counter()-t0)/60:.1f} min")
+
+    # Phase 2: IG sequentially with jitter
+    print(f"\nPhase 2: Instagram (sequential, 2-4s jitter)")
+    t1 = time.perf_counter()
+    consecutive_fails = 0
+    for i, (u, t, l) in enumerate(ig_jobs, 1):
+        res = download_one(u, t, cookies)
+        results.append(res)
+        icon = {"ok": "✓", "skip": "→", "fail": "✗"}.get(res["status"], "?")
+        cat_tag = f"[{res['category']}]" if res.get("category") else ""
+        print(f"  [{i}/{len(ig_jobs)}] {icon} {l:<6} {Path(res['path']).name[:38]:<38} {cat_tag}")
+
+        if res["status"] == "fail" and res.get("category") == "rate_limited":
+            consecutive_fails += 1
+        else:
+            consecutive_fails = 0
+
+        if res["status"] != "skip":
+            if consecutive_fails >= 3:
+                sleep_s = 30 + random.uniform(0, 15)
+                print(f"      ⏸  rate-limit backoff: {sleep_s:.0f}s")
+            else:
+                sleep_s = 2 + random.uniform(0, 2)
+            time.sleep(sleep_s)
+
+    if log_path:
+        Path(log_path).parent.mkdir(parents=True, exist_ok=True)
+        with open(log_path, "w") as f:
+            json.dump(results, f, indent=2)
+
+    # Summary
+    print(f"\n{'═'*60}")
+    print(f"  Total time: {(time.perf_counter()-t0)/60:.1f} min")
+    print(f"  Status: {Counter(r['status'] for r in results)}")
+    for lbl in ["good", "okish", "bad"]:
+        folder = Path(out_root) / lbl
+        n_ok = len(list(folder.glob("*.mp4"))) if folder.exists() else 0
+        print(f"  {lbl:<6} {n_ok:>3} downloaded")
+    print(f"{'═'*60}")
+
+    return results