{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "af42ede4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "import re\n",
    "import time\n",
    "import random\n",
    "import json\n",
    "import tqdm\n",
    "import uuid\n",
    "import requests\n",
    "import string\n",
    "import funcy\n",
    "import gzip\n",
    "import urllib\n",
    "import numpy as np\n",
    "import multiprocessing\n",
    "from bs4 import BeautifulSoup\n",
    "from contextlib import redirect_stderr, redirect_stdout, contextmanager\n",
    "import signal\n",
    "\n",
    "import pandas as pd\n",
    "import cloudscraper\n",
    "import youtube_dl\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.audio.conversion import get_audio_properties, get_duration_s\n",
    "from suno_utils.utils.text import make_unique_list\n",
    "\n",
    "\n",
    "def _courtesy_sleep(avg_sleep_dur_s=0.5):\n",
    "    time.sleep((0.5 + random.random() / 2) * avg_sleep_dur_s)\n",
    "\n",
    "\n",
    "class TimeoutException(Exception): \n",
    "    pass\n",
    "\n",
    "\n",
    "@contextmanager\n",
    "def time_limit(seconds):\n",
    "    seconds = int(round(seconds))\n",
    "    def signal_handler(signum, frame):\n",
    "        raise TimeoutException(\"Function call timed out!\")\n",
    "    signal.signal(signal.SIGALRM, signal_handler)\n",
    "    signal.alarm(seconds)\n",
    "    try:\n",
    "        yield\n",
    "    finally:\n",
    "        signal.alarm(0)\n",
    "        \n",
    "    \n",
    "class Logger():\n",
    "    def __init__(self, filepath):\n",
    "        self._filepath = filepath\n",
    "        self._reset_log()\n",
    "        \n",
    "    def _reset_log(self):\n",
    "        with open(self._filepath, \"w\") as f:\n",
    "            f.write(\"\")\n",
    "        \n",
    "    def _add_line(self, line):\n",
    "        with open(self._filepath, \"a\") as f:\n",
    "            f.write(line + \"\\n\")\n",
    "\n",
    "def _parse_duration(s):\n",
    "    if s is None:\n",
    "        return 0\n",
    "    parts = s.split(\":\")\n",
    "    if len(parts) == 3:\n",
    "        nh, nm, ns = parts\n",
    "    elif len(parts) == 2:\n",
    "        nm, ns = parts\n",
    "        nh = 0\n",
    "    else:\n",
    "        raise ValueError(\"\")\n",
    "    return int(nh) * 60**2 + int(nm) * 60 + int(ns)\n",
    "    \n",
    "\n",
    "def mp_scrape(\n",
    "    extract_f, \n",
    "    queue_items, \n",
    "    global_info=None,\n",
    "    result_filepath=None, \n",
    "    log_filepath=None, \n",
    "    n_cores=5, \n",
    "    chunksize=500, \n",
    "    n_retries=3,\n",
    "    append_results=False,\n",
    "    backoff_dur_s=1.0,\n",
    "    inner_chunksize=1, \n",
    "    quiet=False,\n",
    "):\n",
    "    if global_info is not None:\n",
    "        _f = funcy.partial(extract_f, global_info=global_info)\n",
    "    else:\n",
    "        _f = extract_f\n",
    "    if result_filepath is not None and not append_results:\n",
    "        with open(result_filepath, \"w\") as f:\n",
    "            f.write(\"\")\n",
    "    if log_filepath is not None:\n",
    "        logger = Logger(log_filepath)\n",
    "    out = []\n",
    "    n_chunks = int(np.ceil(len(queue_items) / chunksize))\n",
    "    for n_chunk, queue_items_chunk in tqdm.tqdm(\n",
    "        enumerate(funcy.chunks(chunksize, queue_items)), \n",
    "        total=n_chunks,\n",
    "        disable=quiet,\n",
    "    ):\n",
    "        p = multiprocessing.Pool(n_cores)\n",
    "        t0 = time.time()\n",
    "        remaining_items = [(idx, queue_item) for idx, queue_item in enumerate(queue_items_chunk)]\n",
    "        out_chunk = [None] * len(queue_items_chunk)\n",
    "        for n_retry in range(n_retries):\n",
    "            tmp_out = p.map(_f, [queue_item for _, queue_item in remaining_items], chunksize=inner_chunksize)\n",
    "            _courtesy_sleep(avg_sleep_dur_s=backoff_dur_s)\n",
    "            tmp_remaining_items = []\n",
    "            for (idx, queue_item), tmp_out_item in zip(remaining_items, tmp_out):\n",
    "                out_chunk[idx] = tmp_out_item\n",
    "                if (\n",
    "                    isinstance(tmp_out_item, dict) and (\n",
    "                        (tmp_out_item.get(\"retry\")) == True or \n",
    "                        (\"retry\" not in tmp_out_item and tmp_out_item.get(\"success\") == False)\n",
    "                    )\n",
    "                ):\n",
    "                    tmp_remaining_items.append((idx, queue_item))\n",
    "                    continue\n",
    "            remaining_items = tmp_remaining_items\n",
    "            if len(remaining_items) == 0:\n",
    "                break\n",
    "            if n_retry < n_retries - 1:\n",
    "                logger._add_line(f\"  retrying for {len(remaining_items)}/{len(out_chunk)} items\")\n",
    "            \n",
    "        # show how many failed\n",
    "        n_failed = len([\n",
    "            e for e in out_chunk if e is None or (isinstance(e, dict) and e.get(\"success\") == False)\n",
    "        ])\n",
    "        logger._add_line(f\"  failed on {n_failed}/{len(out_chunk)} items\")\n",
    "            \n",
    "        # break if none were successful\n",
    "        if all([e is None or (isinstance(e, dict) and e.get(\"success\") == False) for e in out_chunk]):\n",
    "            logger._add_line(f\"{len(out_chunk)}/{len(out_chunk)} items failed, aborting.\")\n",
    "            \n",
    "        if result_filepath is not None:\n",
    "            with open(result_filepath, \"a\") as f:\n",
    "                for e in out_chunk:\n",
    "                    f.write(json.dumps(e) + \"\\n\")\n",
    "        else:\n",
    "            out.extend(out_chunk)\n",
    "        td = time.time() - t0\n",
    "        if log_filepath is not None:\n",
    "            logger._add_line(f\"{n_chunk+1}/{n_chunks} - last step took {round(td / 60, 1)} mins\")\n",
    "        p.close()\n",
    "        p.join()\n",
    "    logger._add_line(f\"done!\")\n",
    "    if result_filepath is not None:\n",
    "        return None\n",
    "    return out\n",
    "\n",
    "US_PROXY = (\n",
    "     \"http://brd-customer-hl_98887cab-zone-us_proxy-route_err-block-country-us:\" +\n",
    "     \"3if5he8elfe7@zproxy.lum-superproxy.io:22225\"\n",
    ")\n",
    "\n",
    "DATA_DIR = \"/data2/suno/data/harvest/youtube\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abf21eb5",
   "metadata": {},
   "source": [
    "## Download youtube audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ba83bd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # videos_by_domain_global.jsonl\n",
    "# 12620995 videos found\n",
    "# 6831596 unique\n",
    "# 1668065.9 hours of data\n",
    "# 1339629.9 hours of data with clipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eb36652c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000000 already seen\n"
     ]
    }
   ],
   "source": [
    "# get already processed IDs\n",
    "processed_ids = set()\n",
    "with open(os.path.join(DATA_DIR, \"youtube_metas.jsonl\")) as f:\n",
    "    for line in f:\n",
    "        line = line.strip()\n",
    "        if len(line) == 0:\n",
    "            continue\n",
    "        m = json.loads(line)\n",
    "        processed_ids.add(m[\"id\"])\n",
    "with open(os.path.join(DATA_DIR, \"youtube_metas_p2.jsonl\")) as f:\n",
    "    for line in f:\n",
    "        line = line.strip()\n",
    "        if len(line) == 0:\n",
    "            continue\n",
    "        m = json.loads(line)\n",
    "        processed_ids.add(m[\"id\"])\n",
    "print(len(processed_ids), \"already seen\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d49a97a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▋                                           | 100000/6771507 [00:48<54:02, 2057.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5525564 remaining unique videos\n",
      "306032 excluded,  1 failed\n",
      "861276.3 hours of data\n"
     ]
    }
   ],
   "source": [
    "tot_n = 0\n",
    "durations_s = []\n",
    "youtube_ids = []\n",
    "seen_ids = set()\n",
    "excluded_ids = set()\n",
    "failed_ids = set()\n",
    "with open(os.path.join(DATA_DIR, \"videos_by_domain_global.jsonl\")) as f:\n",
    "    for l in tqdm.tqdm(f, total=100_000):\n",
    "        if len(l.strip()) == 0:\n",
    "            continue\n",
    "        m = json.loads(l)\n",
    "        if m is None:\n",
    "            continue\n",
    "        domain_str, search_results = m\n",
    "        for search_result in search_results:\n",
    "            _id = search_result[\"id\"]\n",
    "            if _id not in seen_ids and _id not in processed_ids:\n",
    "                seen_ids.add(_id)\n",
    "                try:\n",
    "                    ds = _parse_duration(search_result[\"duration\"])\n",
    "                    if ds < 10 or ds > 1 * 60 * 60:\n",
    "                        excluded_ids.add(_id)\n",
    "                        continue\n",
    "                except:\n",
    "                    failed_ids.add(_id)\n",
    "                    continue\n",
    "                durations_s.append(ds)\n",
    "                youtube_ids.append(_id)\n",
    "            tot_n += 1\n",
    "            \n",
    "# print(tot_n, \"videos found\")\n",
    "print(len(youtube_ids), \"remaining unique videos\")\n",
    "print(len(excluded_ids), \"excluded, \", len(failed_ids), \"failed\")\n",
    "print(round(np.sum(durations_s) / 60 / 60, 1), \"hours of data\")\n",
    "\n",
    "youtube_ids = list(youtube_ids)\n",
    "random.seed(6006)\n",
    "random.shuffle(youtube_ids)\n",
    "# 6771507 unique videos\n",
    "# 60091 excluded\n",
    "# 1 failed\n",
    "# 1447996.6 hours of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3138eecd",
   "metadata": {},
   "outputs": [],
   "source": [
    "YDL_OPTS = {\n",
    "    \"format\": \"worstaudio\",\n",
    "    \"outtmpl\": os.path.join(DATA_DIR, \"audio\", \"%(id)s.%(ext)s\"),\n",
    "    \"writesubtitles\": True,\n",
    "    \"allsubtitles\": True,\n",
    "    \"subtitlesformat\": \"best\",\n",
    "    \"socket_timeout\": 5.0,\n",
    "    \"proxy\": US_PROXY,\n",
    "}\n",
    "MAX_DOWNLOAD_TIME_S = 60 * 2\n",
    "\n",
    "YOUTUBE_BASE_URL = \"https://www.youtube.com/watch?v=\"\n",
    "\n",
    "def _resolve_youtube(youtube_id):\n",
    "    url = YOUTUBE_BASE_URL + youtube_id\n",
    "    out_data = {\n",
    "        \"id\": youtube_id,\n",
    "    }\n",
    "    t0 = time.time()\n",
    "    try:\n",
    "        with redirect_stderr(open(os.devnull, \"w\")):\n",
    "            with redirect_stdout(open(os.devnull, \"w\")):\n",
    "                with youtube_dl.YoutubeDL(YDL_OPTS) as ydl:\n",
    "                    # pre-download to check if we want it (duration, views etc)\n",
    "#                     info = ydl.extract_info(url, download=False)\n",
    "                    if MAX_DOWNLOAD_TIME_S is not None:\n",
    "                        with time_limit(MAX_DOWNLOAD_TIME_S):\n",
    "                            info = ydl.extract_info(url, download=True)\n",
    "                    else:\n",
    "                        info = ydl.extract_info(url, download=True)\n",
    "        _ = info.pop(\"formats\", None)\n",
    "        _ = info.pop(\"thumbnails\", None)\n",
    "        _ = info.pop(\"thumbnail\", None)\n",
    "        _ = info.pop(\"automatic_captions\", None)\n",
    "        out_data[\"success\"] = True\n",
    "        out_data[\"retry\"] = False\n",
    "        out_data[\"meta\"] = info\n",
    "        out_data[\"audio_filename\"] = f\"{info['id']}.{info['ext']}\"\n",
    "    except Exception as e:\n",
    "        out_data[\"success\"] = False\n",
    "        out_data[\"retry\"] = (\n",
    "            \"TimeoutException\" in str(type(e)) or\n",
    "            \"unable to download video data\" in str(e) or \n",
    "            \"No video formats found\" in str(e)\n",
    "        )\n",
    "        out_data[\"fail_type\"] = str(type(e))\n",
    "        out_data[\"fail_message\"] = str(e)\n",
    "    t1 = time.time()\n",
    "    out_data[\"runtime_s\"] = round(t1 - t0, 1)\n",
    "    return out_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "edb937ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wc -l /data2/suno/data/harvest/youtube/youtube_metas.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87d9217d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|▏                                             | 12/3000 [1:50:34<458:44:55, 552.71s/it]"
     ]
    }
   ],
   "source": [
    "# ~9.5mins for 1_000 videos, ~20 days for 3M, 100 threads / 32 cores - ~8Tb\n",
    "# bright data limit with 100 servers is 10 Tb per month\n",
    "# started Sun, Nov 13th\n",
    "_ = mp_scrape(\n",
    "    _resolve_youtube, \n",
    "#     youtube_ids[:500_000], # ~2 Tb\n",
    "#     youtube_ids[500_000:1_000_000], # ~2 Tb\n",
    "    youtube_ids[1_000_000:4_000_000], # ~8 Tb\n",
    "    chunksize=1000,\n",
    "    n_cores=100,\n",
    "    n_retries=3,\n",
    "#     append_results=True,\n",
    "    result_filepath=os.path.join(DATA_DIR, \"youtube_metas_p3.jsonl\"),\n",
    "    log_filepath=os.path.join(DATA_DIR, \"logs\", \"audio_harvest.log\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e68e2191",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: do we need to reorganize audio dir to not get too large?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "af2ce70e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# !ls -1 /data2/suno/data/harvest/youtube/audio | grep -v vtt | wc -l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59ea1a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !ffprobe -v error -select_streams a:0 -show_entries stream=codec_name -of default=nokey=1:noprint_wrappers=1 test.webm\n",
    "# !ffmpeg -i test.webm -vn -map 0:a -acodec copy test.opus  # 0:a smooshes, 0:a:0 takes first audio\n",
    "# looks like all streams have just one audio track. webm is opus and m4a is aac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91c6bc01",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc581dc6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "504902b8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a6ca8c0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}