{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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 = \"/data/suno/data/harvest/youtube_med\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abf21eb5",
   "metadata": {},
   "source": [
    "## Download youtube audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d49a97a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "66453it [00:05, 12430.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "465865 unique videos\n",
      "29507 excluded\n",
      "0 failed\n",
      "78053.8 hours of data\n"
     ]
    }
   ],
   "source": [
    "tot_n = 0\n",
    "durations_s = []\n",
    "seen_ids = set()\n",
    "youtube_ids = []\n",
    "excluded_ids = set()\n",
    "failed_ids = set()\n",
    "with open(os.path.join(DATA_DIR, \"videos_urls.jsonl\")) as f:\n",
    "    for l in tqdm.tqdm(f):\n",
    "        if len(l.strip()) == 0:\n",
    "            continue\n",
    "        m = json.loads(l)\n",
    "        if m is None:\n",
    "            continue\n",
    "        if not m[\"success\"]:\n",
    "            continue\n",
    "        for mm in m[\"search_results\"]:\n",
    "            _id = mm[\"id\"]\n",
    "            if _id not in seen_ids:\n",
    "                seen_ids.add(_id)\n",
    "                try:\n",
    "                    ds = _parse_duration(mm[\"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",
    "                seen_ids.add(_id)\n",
    "            tot_n += 1\n",
    "            \n",
    "# print(tot_n, \"videos found\")\n",
    "print(len(youtube_ids), \"unique videos\")\n",
    "print(len(excluded_ids), \"excluded\")\n",
    "print(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",
    "# 490851 unique videos\n",
    "# 4521 excluded\n",
    "# 0 failed\n",
    "# 114425.6 hours of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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",
    "        # TODO: figure out language info somehow\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": 4,
   "id": "edb937ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wc -l /data/suno/data/harvest/youtube_lg/youtube_metas.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "87d9217d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 466/466 [73:14:52<00:00, 565.86s/it]\n"
     ]
    }
   ],
   "source": [
    "# ~7mins for 500 videos, ~4 days for 300k (~100k hours) at 100 cores - ~4Tb\n",
    "# ~24h - 750Gb - 35k hours (from ~200k IDs input)\n",
    "# bright data limit with 100 servers is 10 Tb per month\n",
    "_ = mp_scrape(\n",
    "    _resolve_youtube, \n",
    "    youtube_ids, \n",
    "    chunksize=1000,\n",
    "    n_cores=100,\n",
    "    n_retries=3,\n",
    "#     append_results=True,\n",
    "    result_filepath=os.path.join(DATA_DIR, \"audio_metas.jsonl\"),\n",
    "    log_filepath=os.path.join(DATA_DIR, \"logs\", \"audio_harvest.log\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af2ce70e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# !ls -1 /data/suno/data/harvest/youtube_med/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. wemb is opus and m4a is aac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99fa5918",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: HjUBP5l_2Ys is spanish with english subtitles...\n",
    "#  figure out how to get video languages!!\n",
    "#   https://stackoverflow.com/questions/47105402/is-it-possible-to-obtain-the-spoken-language-from-youtube-dl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3a5d7ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: YOgIA4Drgi0 says human but looks auto gen-ed??? Sentel vs Zentel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43f3de79",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3abcdb5f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d328c9c2",
   "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
}
