{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "204eb465",
   "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": "code",
   "execution_count": 4,
   "id": "ffac48cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "SUBTITLES_DIR = os.path.join(DATA_DIR, \"subtitles\")\n",
    "AUDIO_DIR = os.path.join(DATA_DIR, \"audio\")\n",
    "LOGS_DIR = os.path.join(DATA_DIR, \"logs\")\n",
    "\n",
    "os.makedirs(DATA_DIR, exist_ok=True)\n",
    "os.makedirs(SUBTITLES_DIR, exist_ok=True)\n",
    "os.makedirs(AUDIO_DIR, exist_ok=True)\n",
    "os.makedirs(LOGS_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b64248e1",
   "metadata": {},
   "source": [
    "### Use youtube FE search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "acce1ff8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "66453 medical seed terms\n"
     ]
    }
   ],
   "source": [
    "with open(\"seed_terms.json\") as f:\n",
    "    med_seed_terms = json.load(f)\n",
    "print(len(med_seed_terms), \"medical seed terms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38389aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we want ~300k videos for 50k hours\n",
    "# search gives max 500 results per term, 20 per page\n",
    "# with 50k seed terms we want ~1 pages per term"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c3a50966",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "import re\n",
    "import time\n",
    "import random\n",
    "\n",
    "def _parse_youtube_search_results(html_str):\n",
    "    start_idx = html_str.index(\"ytInitialData\") + len(\"ytInitialData\") + 3\n",
    "    end_idx = html_str.index(\"};\", start_idx) + 1\n",
    "    json_str = html_str[start_idx:end_idx]\n",
    "    page_info = json.loads(json_str)\n",
    "    primary_content = page_info[\"contents\"][\"twoColumnSearchResultsRenderer\"].get(\"primaryContents\")\n",
    "    if primary_content is not None:\n",
    "        content_blocks = primary_content[\"sectionListRenderer\"][\"contents\"]\n",
    "    else:\n",
    "        coninuation_content = page_info[\"onResponseReceivedCommands\"][0][\"appendContinuationItemsAction\"]\n",
    "        content_blocks = coninuation_content[\"continuationItems\"]\n",
    "    results = []\n",
    "    for contents in content_blocks[:-1]:\n",
    "        for video in contents[\"itemSectionRenderer\"][\"contents\"]:\n",
    "            res = {}\n",
    "            if \"videoRenderer\" in video.keys():\n",
    "                video_data = video.get(\"videoRenderer\", {})\n",
    "                res[\"id\"] = video_data.get(\"videoId\", None)\n",
    "                res[\"title\"] = video_data.get(\"title\", {}).get(\"runs\", [[{}]])[0].get(\"text\", None)\n",
    "                res[\"long_desc\"] = video_data.get(\"descriptionSnippet\", {}).get(\"runs\", [{}])[0].get(\"text\", None)\n",
    "                res[\"channel\"] = video_data.get(\"longBylineText\", {}).get(\"runs\", [[{}]])[0].get(\"text\", None)\n",
    "                res[\"duration\"] = video_data.get(\"lengthText\", {}).get(\"simpleText\", None)\n",
    "                res[\"views\"] = video_data.get(\"viewCountText\", {}).get(\"simpleText\", None)\n",
    "                res[\"publish_time\"] = video_data.get(\"publishedTimeText\", {}).get(\"simpleText\", None)\n",
    "                res[\"url_suffix\"] = video_data.get(\n",
    "                    \"navigationEndpoint\", {}\n",
    "                ).get(\"commandMetadata\", {}).get(\"webCommandMetadata\", {}).get(\"url\", None)\n",
    "                results.append(res)\n",
    "    if \"continuationItemRenderer\" not in content_blocks[-1]:\n",
    "        continuation_token = None\n",
    "    else:\n",
    "        continuation_endpoint = content_blocks[-1][\"continuationItemRenderer\"][\"continuationEndpoint\"]\n",
    "        continuation_token = continuation_endpoint[\"continuationCommand\"][\"token\"]\n",
    "    return continuation_token, results\n",
    "\n",
    "def _get_youtube_search_results(query_term, max_n_pages=None):\n",
    "    # NOTE: this returns max ~500 results\n",
    "    query_term = urllib.parse.urlencode({\"a\": query_term})[2:]\n",
    "    proxies = {\n",
    "       \"http\": US_PROXY,\n",
    "       \"https\": US_PROXY,\n",
    "    }\n",
    "    base_url = \"https://www.youtube.com/results?\"\n",
    "    all_results = []\n",
    "    seen_ids = set()\n",
    "    n_pages = 1\n",
    "    while True:\n",
    "        if n_pages == 1:\n",
    "            harvest_url = f\"{base_url}search_query={query_term}&hl=en&sp=EgIoAQ%253D%253D\"\n",
    "        else:\n",
    "            harvest_url = f\"{base_url}continuation={continuation_token}&ctoken={continuation_token}&hl=en\"\n",
    "        out = requests.get(harvest_url, proxies=proxies)\n",
    "        assert(out.ok)\n",
    "        continuation_token, results = _parse_youtube_search_results(out.text)\n",
    "        for m in results:\n",
    "            video_id = m[\"id\"]\n",
    "            if video_id in seen_ids:\n",
    "                continue\n",
    "            seen_ids.add(video_id)\n",
    "            all_results.append(m)\n",
    "        if continuation_token is None:\n",
    "            break\n",
    "        if max_n_pages is not None and n_pages >= max_n_pages:\n",
    "            break\n",
    "        n_pages += 1\n",
    "        _courtesy_sleep(avg_sleep_dur_s=0.3)\n",
    "    return all_results\n",
    "\n",
    "def get_youtube_search_results(query_term):\n",
    "    out_data = {\n",
    "        \"query_term\": query_term,\n",
    "    }\n",
    "    try:\n",
    "        search_results = _get_youtube_search_results(query_term, max_n_pages=2)\n",
    "        out_data[\"success\"] = True\n",
    "        out_data[\"retry\"] = False\n",
    "        out_data[\"search_results\"] = search_results\n",
    "    except Exception as e:\n",
    "        out_data[\"success\"] = False\n",
    "        out_data[\"retry\"] = True\n",
    "        out_data[\"fail_type\"] = str(type(e))\n",
    "        out_data[\"fail_message\"] = str(e)\n",
    "    return out_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f90ee2ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 133/133 [3:48:09<00:00, 102.93s/it]\n"
     ]
    }
   ],
   "source": [
    "# ~10 mins for 500 terms (10 cores) --> 2 days for all 100k search terms\n",
    "_ = mp_scrape(\n",
    "    get_youtube_search_results, \n",
    "    med_seed_terms, \n",
    "    chunksize=500,\n",
    "    n_cores=10,\n",
    "    n_retries=2,\n",
    "#     continue_partial=True,\n",
    "    result_filepath=os.path.join(DATA_DIR, \"videos_urls.jsonl\"), \n",
    "    log_filepath=os.path.join(LOGS_DIR, \"video_urls_harvest.log\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "18ab984c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "66453it [00:06, 9688.36it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1841950 videos found\n",
      "495372 unique\n",
      "0 failed\n",
      "137133.4 hours of data\n",
      "107146.6 hours of data with clipping\n"
     ]
    }
   ],
   "source": [
    "# youtube_ids = set()\n",
    "# tot_n = 0\n",
    "# durations_s = []\n",
    "# failed_ids = []\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",
    "#             if mm[\"id\"] not in youtube_ids:\n",
    "#                 try:\n",
    "#                     ds = _parse_duration(mm[\"duration\"])\n",
    "#                 except:\n",
    "#                     failed_ids.append(mm[\"id\"])\n",
    "#                     continue\n",
    "#                 durations_s.append(ds)\n",
    "#                 youtube_ids.add(mm[\"id\"])\n",
    "#             tot_n += 1\n",
    "# youtube_ids = list(youtube_ids)\n",
    "# print(tot_n, \"videos found\")\n",
    "# print(len(youtube_ids), \"unique\")\n",
    "# print(len(failed_ids), \"failed\")\n",
    "# print(round(np.sum(durations_s) / 60 / 60, 1), \"hours of data\")\n",
    "# print(round(np.clip(np.array(durations_s), a_max=3600, a_min=0).sum() / 60 / 60, 1), \"hours of data with clipping\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afb97fcb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f3c5458",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40c594ef",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a615ddc0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "046c0412",
   "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
}