{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac6ddc2b",
   "metadata": {},
   "source": [
    "## get all unique vtt files and do parsing into new file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "70e33ca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d72f7a4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning : `load_model` does not return WordVectorModel or SupervisedModel any more, but a `FastText` object which is very similar.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "import fasttext\n",
    "from multiprocessing import Pool\n",
    "\n",
    "from suno_utils.harvest.youtube.constants import USED_LANG_CODES\n",
    "from suno_utils.harvest.youtube.captions import verify_vtt_lang\n",
    "\n",
    "DATA_DIR = \"/data2/suno/data/harvest/youtube\"\n",
    "\n",
    "fasttext_lang_model = fasttext.load_model(\"lid.176.bin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b765c11a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2280733"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vtt_fns = [fn for fn in os.listdir(os.path.join(DATA_DIR, \"audio\")) if fn.endswith(\"vtt\")]\n",
    "vtt_fns = [fn for fn in vtt_fns if len(fn.split(\".\")) == 3]\n",
    "id_list = [fn.split(\".\")[0] for fn in vtt_fns]\n",
    "s = pd.Series(id_list).value_counts()\n",
    "unique_ids = set(s[s==1].index)\n",
    "filtered_vtt_fns = [\n",
    "    vtt_fn \n",
    "    for vtt_fn in vtt_fns \n",
    "    if (\n",
    "        len(vtt_fn.split(\".\")[1]) >= 2 and \n",
    "        vtt_fn.split(\".\")[0] in unique_ids\n",
    "    )\n",
    "]\n",
    "len(filtered_vtt_fns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "631a0cf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def foo(vtt_fn):\n",
    "    with open(os.path.join(DATA_DIR, \"audio\", vtt_fn)) as f:\n",
    "        vtt_str = f.read()\n",
    "    try:\n",
    "        check_lang_code = vtt_fn.split(\".\")[1]\n",
    "        verify_vtt_lang(vtt_str, check_lang_code, fasttext_lang_model)\n",
    "        verified_vtt_fns.append(vtt_fn)\n",
    "        return vtt_fn\n",
    "    except:\n",
    "        pass\n",
    "    return None\n",
    "\n",
    "with Pool(30) as p:\n",
    "    out = p.map(foo, filtered_vtt_fns, chunksize=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "d3fa84c8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1815144"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "verified_vtt_fns = [e for e in out if e is not None]\n",
    "len(verified_vtt_fns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "1326c974",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"verified_vtt_fns.json\", \"w\") as f:\n",
    "    json.dump(verified_vtt_fns, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a45ea91",
   "metadata": {},
   "source": [
    "## Find tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "502e99f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from suno_utils.utils.text import normalize_whitespace\n",
    "\n",
    "def foo(vtt_fn):\n",
    "    base_lang = vtt_fn.split(\".\")[1][:2]\n",
    "    with open(os.path.join(DATA_DIR, \"audio\", vtt_fn)) as f:\n",
    "        vtt_str = f.read()\n",
    "    try:\n",
    "        _, annotated_lines = parse_vtt(vtt_str)\n",
    "        text = \"\".join([s for _, s in annotated_lines])\n",
    "        tags_br = [\n",
    "            normalize_whitespace(s.lower()) \n",
    "            for s in re.findall(r\"\\[(.*?)\\]\", text) \n",
    "            if len(s.strip()) > 0\n",
    "        ]\n",
    "        tags_pa = [\n",
    "            normalize_whitespace(s.lower()) \n",
    "            for s in re.findall(r\"\\((.*?)\\)\", text) \n",
    "            if len(s.strip()) > 0\n",
    "        ]\n",
    "        if len(tags_br) + len(tags_pa) == 0:\n",
    "            return None\n",
    "        return base_lang, list(set(tags_br)), list(set(tags_pa))\n",
    "    except:\n",
    "        pass\n",
    "    return None\n",
    "\n",
    "with Pool(30) as p:\n",
    "    out = p.map(foo, verified_vtt_fns, chunksize=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "79f447ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1304"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# keep up 100 per language\n",
    "from collections import defaultdict, Counter\n",
    "tags_info = defaultdict(list)\n",
    "for e in out:\n",
    "    if e is None:\n",
    "        continue\n",
    "    lang, a, b = e\n",
    "    tags_info[lang].extend(a)\n",
    "    tags_info[lang].extend(b)\n",
    "retained_tags = []\n",
    "for k, v in tags_info.items():\n",
    "    for kk, vv in Counter(v).most_common(100):\n",
    "        if vv >= 10:\n",
    "            retained_tags.append(kk)\n",
    "len(retained_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "1d9eb360",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import json\n",
    "# print(json.dumps(retained_tags, indent=4, ensure_ascii=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf139e68",
   "metadata": {},
   "source": [
    "## do parsing/cleaning\n",
    "## do filtering base on text (cap etc)\n",
    "## segmenting\n",
    "## extra filter for asr english\n",
    "### back with english-only asr preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb97fca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db5c20f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64166603",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03fc9cf8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b86a43",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d4a62d4c",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddd22c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "vtt_fns = os.listdir(os.path.join(DATA_DIR, \"audio\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ee50a7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "l = [fn.split(\".\")[0] for fn in vtt_fns]\n",
    "vc = pd.Series([e for e in l]).value_counts()\n",
    "unique_ids = set(vc[vc == 1].index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "97d9c1c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "langs = [fn.split(\".\")[1] for fn in vtt_fns if fn.split(\".\")[0] in unique_ids]\n",
    "vc = pd.Series(langs).value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "6194bfa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# can we tell \"KIgOUc63dRo\" is russian?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "id": "155a8129",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[youtube] rKYaV-vRc7E: Downloading webpage\n",
      "[download] Destination: tmp/rKYaV-vRc7E.webm\n",
      "[download] 100% of 1.28MiB in 00:0091MiB/s ETA 00:001\n"
     ]
    }
   ],
   "source": [
    "import youtube_dl\n",
    "\n",
    "YDL_OPTS = {\n",
    "    \"format\": \"worstaudio\",\n",
    "    \"outtmpl\": \"tmp/%(id)s.%(ext)s\",\n",
    "    \"writesubtitles\": True,\n",
    "    \"allsubtitles\": True,\n",
    "    \"subtitlesformat\": \"best\",\n",
    "    \"socket_timeout\": 5.0,\n",
    "}\n",
    "\n",
    "YOUTUBE_BASE_URL = \"https://www.youtube.com/watch?v=\"\n",
    "\n",
    "# url = YOUTUBE_BASE_URL + \"KIgOUc63dRo\"  # russian, all sorts of stuff\n",
    "\n",
    "# url = YOUTUBE_BASE_URL + \"n9pMNI-tWzE\"  # cnn, english only\n",
    "# len(info[\"subtitles\"].keys()) == 1\n",
    "\n",
    "url = YOUTUBE_BASE_URL + \"rKYaV-vRc7E\"  # last week tonight, auto english only\n",
    "\n",
    "with youtube_dl.YoutubeDL(YDL_OPTS) as ydl:\n",
    "    info = ydl.extract_info(url, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f59dd5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: next look at how many have ALL vs just a few in 'manual' category"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "575ac7ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 176,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(info[\"subtitles\"].keys()) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "fcc92b48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['af', 'ak', 'sq', 'am', 'ar', 'hy', 'as', 'ay', 'az', 'bn', 'eu', 'be', 'bho', 'bs', 'bg', 'my', 'ca', 'ceb', 'zh-Hans', 'zh-Hant', 'co', 'hr', 'cs', 'da', 'dv', 'nl', 'en', 'eo', 'et', 'ee', 'fil', 'fi', 'fr', 'gl', 'lg', 'ka', 'de', 'el', 'gn', 'gu', 'ht', 'ha', 'haw', 'iw', 'hi', 'hmn', 'hu', 'is', 'ig', 'id', 'ga', 'it', 'ja', 'jv', 'kn', 'kk', 'km', 'rw', 'ko', 'kri', 'ku', 'ky', 'lo', 'la', 'lv', 'ln', 'lt', 'lb', 'mk', 'mg', 'ms', 'ml', 'mt', 'mi', 'mr', 'mn', 'ne', 'nso', 'no', 'ny', 'or', 'om', 'ps', 'fa', 'pl', 'pt', 'pa', 'qu', 'ro', 'ru', 'sm', 'sa', 'gd', 'sr', 'sn', 'sd', 'si', 'sk', 'sl', 'so', 'st', 'es', 'su', 'sw', 'sv', 'tg', 'ta', 'tt', 'te', 'th', 'ti', 'ts', 'tr', 'tk', 'uk', 'und', 'ur', 'ug', 'uz', 'vi', 'cy', 'fy', 'xh', 'yi', 'yo', 'zu'])"
      ]
     },
     "execution_count": 152,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "info[\"automatic_captions\"].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b16a8eb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "727b3efc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9966d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e7f596",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abf73b28",
   "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
}