{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4a2f4e76",
   "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 html\n",
    "import requests\n",
    "import string\n",
    "import funcy\n",
    "import numpy as np\n",
    "import gzip\n",
    "import multiprocessing\n",
    "from bs4 import BeautifulSoup\n",
    "from contextlib import redirect_stdout, redirect_stderr\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\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 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",
    "    \n",
    "def get_cloudscraper(test_url=\"https://genius.com/songs/all\"):\n",
    "    for _ in range(10):\n",
    "        scraper = cloudscraper.create_scraper(disableCloudflareV1=True)\n",
    "        out = scraper.get(test_url)\n",
    "        if out.ok:\n",
    "            return scraper\n",
    "        _courtesy_sleep()\n",
    "    raise ValueError(\"could not initiate cloudscraper\")\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",
    "    continue_partial=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 continue_partial:\n",
    "        with open(result_filepath, \"w\") as f:\n",
    "            f.write(\"\")\n",
    "    p = multiprocessing.Pool(n_cores)\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",
    "        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",
    "            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 isinstance(tmp_out_item, dict) and tmp_out_item.get(\"success\") == False:\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(queue_items_chunk)} items\")\n",
    "            else:\n",
    "                logger._add_line(f\"  failed on {len(remaining_items)}/{len(queue_items_chunk)} items\")\n",
    "            _courtesy_sleep(avg_sleep_dur_s=backoff_dur_s)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3166f3d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "GENIUS_DIR = \"/data/suno/data/harvest/genius\"\n",
    "GENIUS_HTML_DIR = os.path.join(GENIUS_DIR, \"raw_html\")\n",
    "GENIUS_AUDIO_DIR = os.path.join(GENIUS_DIR, \"audio\")\n",
    "GENIUS_LOG_DIR = os.path.join(GENIUS_DIR, \"logs\")\n",
    "\n",
    "os.makedirs(GENIUS_DIR, exist_ok=True)\n",
    "os.makedirs(GENIUS_HTML_DIR, exist_ok=True)\n",
    "os.makedirs(GENIUS_AUDIO_DIR, exist_ok=True)\n",
    "os.makedirs(GENIUS_LOG_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58580725",
   "metadata": {},
   "source": [
    "## Get featured song list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "255afa40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def _parse_featured_songs(soup):\n",
    "#     song_data = []\n",
    "#     for e in soup.find_all(\"ul\", attrs={\"class\": \"song_list\"}):\n",
    "#         for ee in e.find_all(\"li\"):\n",
    "#             data_id = ee.get_attribute_list(\"data-id\")[0]\n",
    "#             url = ee.find_all(\"a\")[0].get_attribute_list(\"href\")[0]\n",
    "#             title = ee.find_all(\"a\")[0].get_attribute_list(\"title\")[0]\n",
    "#             song_data.append({\n",
    "#                 \"url\": url,\n",
    "#                 \"title\": title,\n",
    "#                 \"data_id\": data_id,\n",
    "#                 \"uid\": str(uuid.uuid4()),\n",
    "#             })\n",
    "#     return song_data\n",
    "\n",
    "\n",
    "# def get_genius_featured_songs(scraper, from_page=1, to_page=3, avg_sleep_dur_s=0.5):\n",
    "#     assert(to_page <= 100)\n",
    "#     all_songs_data = []\n",
    "#     for n in tqdm.tqdm(range(from_page, to_page + 1)):\n",
    "#         out = scraper.get(f\"https://genius.com/songs/all?page={n}\")\n",
    "#         assert(out.ok)\n",
    "#         soup = BeautifulSoup(out.text)\n",
    "#         all_songs_data.extend(_parse_featured_songs(soup))\n",
    "#         _courtesy_sleep(avg_sleep_dur_s=avg_sleep_dur_s)\n",
    "#     return all_songs_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d4da7024",
   "metadata": {},
   "outputs": [],
   "source": [
    "# featured_song_metas = get_genius_featured_songs(scraper, from_page=1, to_page=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "e95de510",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(os.path.join(GENIUS_DIR, \"featured_song_metas.json\"), \"w\") as f:\n",
    "#     json.dump(featured_song_metas, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7965050",
   "metadata": {},
   "source": [
    "## Get all songs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "d3ffd5e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = Logger(os.path.join(GENIUS_LOG_DIR, \"artists_harvest.log\"))\n",
    "\n",
    "def get_genius_artists(scraper, running_results=None, start_idx=0, avg_sleep_dur_s=0.2):\n",
    "    if running_results is not None:\n",
    "        artist_links = running_results[:]\n",
    "    else:\n",
    "        artist_links = []\n",
    "    for c in tqdm.tqdm((string.ascii_lowercase + \"0\")[start_idx:]):\n",
    "        n_page = 1\n",
    "        while True:\n",
    "            logger._add_line(f\"getting '{c}' page {n_page}.\")\n",
    "            for n in range(3):\n",
    "                out = scraper.get(f\"https://genius.com/artists-index/{c}/all?page={n_page}\")\n",
    "                if out.ok:\n",
    "                    break\n",
    "                logger._add_line(f\"retrying...\")\n",
    "                _courtesy_sleep(avg_sleep_dur_s=avg_sleep_dur_s)\n",
    "            assert(out.ok)\n",
    "            soup = BeautifulSoup(out.text)\n",
    "            tmp_artists = []\n",
    "            for e in soup.find_all(\"ul\", attrs={\"class\": \"artists_index_list\"}):\n",
    "                for ee in e.find_all(\"li\"):\n",
    "                    tmp_artists.append(ee.a.attrs[\"href\"])\n",
    "            if len(tmp_artists) == 0:\n",
    "                break\n",
    "            artist_links.extend(tmp_artists)\n",
    "            n_page += 1\n",
    "            _courtesy_sleep(avg_sleep_dur_s=avg_sleep_dur_s)\n",
    "        with open(os.path.join(GENIUS_DIR, \"artist_links.json\"), \"w\") as f:\n",
    "            json.dump(artist_links, f)\n",
    "    return artist_links"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a983468a",
   "metadata": {},
   "outputs": [],
   "source": [
    "scraper = get_cloudscraper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "855ae9c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████| 24/24 [1:41:16<00:00, 253.20s/it]\n"
     ]
    }
   ],
   "source": [
    "artist_links = get_genius_artists(scraper)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6565965",
   "metadata": {},
   "source": [
    "### resolve artist names to IDs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "989e4fe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(GENIUS_DIR, \"artist_links.json\")) as f:\n",
    "    artist_links = json.load(f)\n",
    "artist_slugs = [url.strip(\" /\").split(\"/\")[-1] for url in artist_links]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb1fdb9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _resolve_artist_meta_1(artist_slug):\n",
    "    try:\n",
    "        out = requests.get(f\"https://genius.com/api/search/artist?q={artist_slug}&per_page=10&page=1\")\n",
    "        assert(out.ok)\n",
    "        out = out.json()\n",
    "        for e in out[\"response\"][\"sections\"][0][\"hits\"]:\n",
    "            meta = e[\"result\"]\n",
    "            if meta[\"slug\"] == artist_slug:\n",
    "                artist_info = {\n",
    "                    \"id\": meta[\"id\"],\n",
    "                    \"is_meme_verified\": meta[\"is_meme_verified\"],\n",
    "                    \"is_verified\": meta[\"is_verified\"],\n",
    "                    \"slug\": meta[\"slug\"],\n",
    "                    \"name\": meta[\"name\"],\n",
    "                }\n",
    "                break\n",
    "        return artist_info\n",
    "    except:\n",
    "        pass\n",
    "    _courtesy_sleep(avg_sleep_dur_s=0.1)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "f95306e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████| 1710/1710 [13:11:43<00:00, 27.78s/it]\n"
     ]
    }
   ],
   "source": [
    "artist_metas = mp_scrape(\n",
    "    _resolve_artist_meta_1, \n",
    "    artist_slugs, \n",
    "    tmp_result_filepath=os.path.join(GENIUS_DIR, \"artist_metas_1.json\"), \n",
    "    log_filepath=os.path.join(GENIUS_LOG_DIR, \"artist_metas_1_harvest.log\"), \n",
    "    n_cores=5, \n",
    "    chunksize=500, \n",
    "    inner_chunksize=1, \n",
    "    quiet=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "563e50a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "40849"
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# resolve Nones with actual artist url\n",
    "unresolved_artist_slugs = []\n",
    "for n, (artist_slug, artist_meta) in enumerate(zip(artist_slugs, artist_metas)):\n",
    "    if artist_meta is None:\n",
    "        unresolved_artist_slugs.append((n, artist_slug))\n",
    "len(unresolved_artist_slugs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "54e51f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _resolve_artist_meta_2(queue_item):\n",
    "    idx, artist_slug = queue_item\n",
    "    try:\n",
    "        out = requests.get(f\"https://genius.com/artists/{artist_slug}\")\n",
    "        assert(out.ok)\n",
    "        m = re.search(r\"\\{\\\"name\\\"\\:\\\"artist\\_id\\\"\\,\\\"values\\\"\\:\\[\\\"([0-9]+)\\\"\\]\", out.text, flags=re.DOTALL)\n",
    "        artist_id = m.group(1)\n",
    "        out = requests.get(f\"https://genius.com/api/artists/{artist_id}\")\n",
    "        assert(out.ok)\n",
    "        out_json = out.json()\n",
    "        meta = out_json[\"response\"][\"artist\"]\n",
    "        artist_info = {\n",
    "            \"id\": artist_id,\n",
    "            \"is_meme_verified\": meta[\"is_meme_verified\"],\n",
    "            \"is_verified\": meta[\"is_verified\"],\n",
    "            \"slug\": artist_slug,\n",
    "            \"name\": meta[\"name\"],\n",
    "        }\n",
    "        return idx, artist_info\n",
    "    except:\n",
    "        pass\n",
    "    _courtesy_sleep(avg_sleep_dur_s=0.1)\n",
    "    return idx, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "id": "33ae38c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████| 82/82 [2:01:13<00:00, 88.70s/it]\n"
     ]
    }
   ],
   "source": [
    "out = mp_scrape(\n",
    "    _resolve_artist_meta_2, \n",
    "    unresolved_artist_slugs, \n",
    "    chunksize=500,\n",
    "    tmp_result_filepath=os.path.join(GENIUS_DIR, \"artist_metas_2.json\"), \n",
    "    log_filepath=os.path.join(GENIUS_LOG_DIR, \"artist_metas_2_harvest.log\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "id": "aa5d91cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "40849/854984 unresolved ones before\n",
      "10/854984 unresolved ones after\n",
      "854974 valid artists\n",
      "25488 verified artists\n"
     ]
    }
   ],
   "source": [
    "print(\"{}/{} unresolved ones before\".format(len([e for e in artist_metas if e is None]), len(artist_metas)))\n",
    "fixed_artist_metas = artist_metas[:]\n",
    "for idx, artist_meta in out:\n",
    "    fixed_artist_metas[idx] = artist_meta\n",
    "print(\"{}/{} unresolved ones after\".format(len([e for e in fixed_artist_metas if e is None]), len(artist_metas)))\n",
    "valid_artist_metas = [m for m in fixed_artist_metas if m is not None]\n",
    "print(len(valid_artist_metas), \"valid artists\")\n",
    "print(len([m for m in valid_artist_metas if m[\"is_verified\"]]), \"verified artists\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "id": "d6fa382a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(GENIUS_DIR, \"artist_metas.json\"), \"w\") as f:\n",
    "    json.dump(valid_artist_metas, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94bb1b9f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8d44557",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "eda55d45",
   "metadata": {},
   "source": [
    "### get song urls for all artists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9c3ab235",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(GENIUS_DIR, \"artist_metas.json\")) as f:\n",
    "    artist_metas = json.load(f)\n",
    "artist_ids = [m[\"id\"] for m in artist_metas]\n",
    "artist_id2slug = {m[\"id\"]: m[\"slug\"] for m in artist_metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fda1ff39",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _resolve_artist_songs(artist_id, global_info):\n",
    "    raw_html_dir = global_info[\"raw_html_dir\"]\n",
    "    os.makedirs(raw_html_dir, exist_ok=True)\n",
    "    out_meta = {\n",
    "        \"artist_id\": artist_id,\n",
    "        \"song_metas\": None,\n",
    "    }\n",
    "    try:\n",
    "        song_metas = []\n",
    "        n_page = 1\n",
    "        while True:\n",
    "            url = f\"https://genius.com/api/artists/{artist_id}/songs?page={n_page}&per_page=50\"\n",
    "            html_fp = os.path.join(raw_html_dir, f\"{artist_id}_p{n_page}.json.gz\")\n",
    "            if os.path.exists(html_fp):\n",
    "                with gzip.open(html_filepath, \"rb\") as f:\n",
    "                    out_json = json.load(f)\n",
    "            else:\n",
    "                for n_retry in range(3):\n",
    "                    out = requests.get(url)\n",
    "                    if not out.ok:\n",
    "                        _courtesy_sleep(avg_sleep_dur_s=0.1)\n",
    "                        continue\n",
    "                assert(out.ok)\n",
    "                out_json = out.json()\n",
    "                with gzip.open(html_fp, \"wb\") as f:\n",
    "                    json.dump(out_json, f)\n",
    "            tmp_song_metas = out_json[\"response\"][\"songs\"]\n",
    "            song_metas.extend([\n",
    "                {\n",
    "                    \"id\": m[\"id\"],\n",
    "                    \"path\": m[\"path\"].strip(\" /\"),\n",
    "                    \"annotation_count\": m[\"annotation_count\"],\n",
    "                    \"title\": m[\"title\"],\n",
    "                    \"language\": m[\"language\"],\n",
    "                    \"lyrics_state\": m[\"lyrics_state\"],\n",
    "                    \"instrumental\": m[\"instrumental\"],\n",
    "                    \"artist_ids\": [e[\"id\"] for e in [m[\"primary_artist\"]] + m[\"featured_artists\"]],\n",
    "                    \"updated_by_human_at\": m[\"updated_by_human_at\"],\n",
    "                } for m in tmp_song_metas\n",
    "            ])\n",
    "            if out_json[\"response\"][\"next_page\"] is None:\n",
    "                break\n",
    "            n_page += 1\n",
    "        out_meta[\"song_metas\"] = song_metas\n",
    "    except:\n",
    "        pass\n",
    "    _courtesy_sleep(avg_sleep_dur_s=0.1)\n",
    "    return out_meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6ee86dde",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████| 855/855 [16:39:04<00:00, 70.11s/it]\n"
     ]
    }
   ],
   "source": [
    "_ = mp_scrape(\n",
    "    _resolve_artist_songs, \n",
    "    artist_ids, \n",
    "    chunksize=1000,\n",
    "    n_cores=10,\n",
    "    global_info={\n",
    "        \"raw_html_dir\": os.path.join(GENIUS_HTML_DIR, \"artist_songs\"),\n",
    "    },\n",
    "    result_filepath=os.path.join(GENIUS_DIR, \"artist_songs.jsonl\"), \n",
    "    log_filepath=os.path.join(GENIUS_LOG_DIR, \"artist_songs_harvest.log\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8e2e65ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "854974/854679 artists parsed, 20 failed, 6392 with no songs\n",
      "15438509 songs, 6178316 unique\n"
     ]
    }
   ],
   "source": [
    "artist_songs = []\n",
    "n_artists = 0\n",
    "n_failed = 0\n",
    "with open(os.path.join(GENIUS_DIR, \"artist_songs.jsonl\")) as f:\n",
    "    for l in f:\n",
    "        if len(l.strip()) == 0:\n",
    "            break\n",
    "        n_artists += 1\n",
    "        d = json.loads(l)\n",
    "        if d[\"song_metas\"] is None:\n",
    "            n_failed += 1\n",
    "            continue\n",
    "        artist_songs.append((artist_id2slug[d[\"artist_id\"]], [m[\"id\"] for m in d[\"song_metas\"]]))\n",
    "\n",
    "print(\n",
    "    \"{}/{} artists parsed, {} failed, {} with no songs\".format(\n",
    "        n_artists,\n",
    "        len(artist_id2slug),\n",
    "        n_failed,\n",
    "        len([1 for _, song_ids in artist_songs if len(song_ids) == 0]),\n",
    "    )\n",
    ")\n",
    "print(\n",
    "    \"{} songs, {} unique\".format(\n",
    "        np.sum([len(song_ids) for _, song_ids in artist_songs]),\n",
    "        len(set(funcy.flatten([song_ids for _, song_ids in artist_songs])))\n",
    "    )\n",
    ")\n",
    "# 854974/854679 artists parsed, 20 failed, 6392 with no songs\n",
    "# 15438509 songs, 6178316 unique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "573d2dba",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cb9b54c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "38c121c7",
   "metadata": {},
   "source": [
    "## Get song details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "07f03d3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# song_id=324459, artist_id=1167, https://genius.com/Johnny-cash-ghost-riders-in-the-sky-lyrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ad54e2ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _parse_song_html(soup):\n",
    "    title = None\n",
    "    artist = None\n",
    "    for e in soup.find_all(\"h1\"):\n",
    "        l = [ee for ee in e.get_attribute_list(\"class\") if ee is not None]\n",
    "        if any([re.search(r\"SongHeader[a-z]+\\_\\_Title\", ee) for ee in l]):\n",
    "            title = e.get_text()\n",
    "    for e in soup.find_all(\"a\"):\n",
    "        l = [ee for ee in e.get_attribute_list(\"class\") if ee is not None]\n",
    "        if any([re.search(r\"SongHeader[a-z]+\\_\\_Artist\", ee) for ee in l]):\n",
    "            artist = e.get_text()\n",
    "    lyrics_containers = soup.find_all(attrs={\"data-lyrics-container\": \"true\"})\n",
    "    tmp_container = []\n",
    "    for lyrics_container in lyrics_containers:\n",
    "        tmp_container.append(lyrics_container.get_text())\n",
    "    lyrics = \"\\n\\n\".join(tmp_container)\n",
    "    return title, artist, lyrics\n",
    "\n",
    "def _unescape(m):\n",
    "    s = m.group()\n",
    "    return s[:-2]\n",
    "\n",
    "def _get_extra_info(html_str):\n",
    "    json_str = re.search(r\"JSON\\.parse\\(\\'(.+?)\\'\\)\\;\", html.unescape(html_str)).group(1)\n",
    "    json_str = json_str.replace(\"\\\\n\", \"<br>\")\n",
    "    json_str = re.sub(r\"\\\\+\", _unescape, json_str)\n",
    "    m = json.loads(json_str)\n",
    "    # get tags & pageviews\n",
    "    details_key = None\n",
    "    for k, v in m[\"songPage\"].items():\n",
    "        if isinstance(v, list) and any([isinstance(vv, dict) and vv.get(\"name\", \"\") == \"song_id\" for vv in v]):\n",
    "            details_key = k\n",
    "            break\n",
    "    info = m[\"songPage\"][details_key]\n",
    "    page_views = None\n",
    "    for d in info:\n",
    "        if \"name\" in d and d[\"name\"] == \"pageviews\":\n",
    "            page_views = int(d[\"values\"][0])\n",
    "            break\n",
    "    tags = None\n",
    "    for d in info:\n",
    "        if \"name\" in d and d[\"name\"] == \"tag_id\":\n",
    "            tags = [int(e) for e in d[\"values\"]]\n",
    "            break\n",
    "    # get more song details\n",
    "    song_id = str(m[\"songPage\"][\"song\"])\n",
    "    song_details = m[\"entities\"][\"songs\"][song_id]\n",
    "    youtube_start = song_details[\"youtubeStart\"]\n",
    "    if youtube_start is None or len(youtube_start) == 0:\n",
    "        youtube_start = \"0\"\n",
    "    song_info = {\n",
    "        \"id\": song_id,\n",
    "        \"youtube_url\": song_details[\"youtubeUrl\"],\n",
    "        \"youtube_start\": youtube_start,\n",
    "        \"lang\": song_details[\"language\"],\n",
    "        \"views\": page_views,\n",
    "        \"tags\": tags,\n",
    "    }\n",
    "    return song_info\n",
    "\n",
    "def _parse_song_details(html_str):\n",
    "    # parse basic data\n",
    "    song_info = _get_extra_info(html_str)\n",
    "    # parse some more data\n",
    "    soup = BeautifulSoup(re.sub(r\"\\<\\s*br\\s*\\/?\\s*\\>\", \"\\n\", html_str))\n",
    "    title, artist, lyrics = _parse_song_html(soup)\n",
    "    song_info[\"title\"] = title\n",
    "    song_info[\"artist\"] = artist\n",
    "    song_info[\"lyrics\"] = lyrics\n",
    "    youtube_urls = re.findall(\n",
    "        r\"[Uu]rl\\\\\\\"\\:\\\\\\\"(https?\\:\\/\\/www\\.youtube\\.com\\/watch\\?v\\=[a-zA-Z0-9\\_\\-]{11})\", \n",
    "        html_str, \n",
    "        flags=re.DOTALL,\n",
    "    )\n",
    "    if len(youtube_urls) > 0:\n",
    "        youtube_url = youtube_urls[0]\n",
    "    else:\n",
    "        youtube_url = None\n",
    "    song_info[\"youtube_url_2\"] = youtube_url\n",
    "    return song_info\n",
    "\n",
    "\n",
    "def get_song_details(song_slug, global_info):\n",
    "    raw_html_dir = global_info[\"raw_html_dir\"]\n",
    "    os.makedirs(raw_html_dir, exist_ok=True)\n",
    "    try:\n",
    "        url = f\"https://genius.com/{song_slug}\"\n",
    "        html_filepath = os.path.join(raw_html_dir, f\"{song_slug}.txt.gz\")\n",
    "        if os.path.exists(html_filepath):\n",
    "            with gzip.open(html_filepath, \"rb\") as f:\n",
    "                html_str = f.read().decode()\n",
    "        else:\n",
    "            raise ValueError(\"skip\")\n",
    "            out = requests.get(url, timeout=0.5)\n",
    "            assert(out.ok)\n",
    "            _courtesy_sleep(avg_sleep_dur_s=0.2)\n",
    "            html_str = out.text\n",
    "            with gzip.open(html_filepath, \"wb\") as f:\n",
    "                f.write(html_str.encode())\n",
    "        song_info = _parse_song_details(html_str)\n",
    "        out_meta = {\n",
    "            \"success\": True,\n",
    "            \"song_slug\": song_slug,\n",
    "            \"meta\": song_info,\n",
    "        }\n",
    "    except Exception as e:\n",
    "        out_meta = {\n",
    "            \"success\": False,\n",
    "            \"song_slug\": song_slug,\n",
    "            \"fail_type\": str(type(e)),\n",
    "            \"fail_message\": str(e),\n",
    "        }\n",
    "    return out_meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c56a017a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: can also get youtube and other info here: https://genius.com/api/songs/2946673"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "594c8917",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "854974it [01:13, 11604.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5480217 song URLs\n"
     ]
    }
   ],
   "source": [
    "song_slugs = []\n",
    "with open(os.path.join(GENIUS_DIR, \"artist_songs.jsonl\")) as f:\n",
    "    for l in tqdm.tqdm(f):\n",
    "        if len(l.strip()) == 0:\n",
    "            break\n",
    "        d = json.loads(l)\n",
    "        if d[\"song_metas\"] is None:\n",
    "            continue\n",
    "        for song_meta in d[\"song_metas\"]:\n",
    "            if (\n",
    "#                 song_meta[\"language\"] == \"en\" and \n",
    "                song_meta[\"lyrics_state\"] == \"complete\" and \n",
    "                song_meta[\"instrumental\"] == False\n",
    "            ):\n",
    "                song_slugs.append(song_meta[\"path\"])\n",
    "song_slugs = sorted(list(set(song_slugs)))\n",
    "print(len(song_slugs), \"song URLs\")\n",
    "# 5480217 total\n",
    "# 3687035 en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "aedafeb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# success_slugs = set()\n",
    "# with open(os.path.join(GENIUS_DIR, \"song_details.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",
    "#         if m[\"success\"]:  # TODO: or type: not available or something to not redo obvious failures\n",
    "#             success_slugs.add(m[\"song_slug\"])\n",
    "\n",
    "# remaining_song_slugs = sorted(list(set(song_slugs) - success_slugs))\n",
    "# print(len(remaining_song_slugs), \"songs remaining\")\n",
    "# # 3491355 songs remaining\n",
    "# # 3222377 songs remaining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "bf463552",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1097/1097 [3:07:37<00:00, 10.26s/it]\n"
     ]
    }
   ],
   "source": [
    "# ~30 mins for 10_000 songs (10 cores) --> 14 days for all 6 million songs\n",
    "_ = mp_scrape(\n",
    "    get_song_details, \n",
    "    song_slugs, \n",
    "    chunksize=5000,\n",
    "    n_cores=20,  # 5\n",
    "    n_retries=1, # 5,\n",
    "    continue_partial=False, #True,\n",
    "    global_info={\n",
    "        \"raw_html_dir\": os.path.join(GENIUS_HTML_DIR, \"song_details\"),\n",
    "    },\n",
    "    result_filepath=os.path.join(GENIUS_DIR, \"song_details.jsonl\"), \n",
    "    log_filepath=os.path.join(GENIUS_LOG_DIR, \"song_details_harvest.log\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78480525",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389f0a2e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe0d63af",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0dfae2d6",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dedc1e9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef860a3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5abf5fa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f706594a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5a21041",
   "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
}