{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "938faf74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run scraping\n",
    "\n",
    "# /home/georg/venvs/ml/bin/python -m suno_utils.scripts.supreme_court.fetch_data \\\n",
    "#     --raw_data_dir=/mnt/data-ssd-1/data/supreme_corpus/raw_data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d7b9363c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Run parsing\n",
    "\n",
    "# /home/georg/venvs/ml/bin/python -m suno_utils.scripts.supreme_court.process_data \\\n",
    "#     --raw_data_dir=/mnt/data-ssd-1/data/supreme_corpus/raw_data/ \\\n",
    "#     --processed_data_dir=/mnt/data-ssd-1/data/supreme_corpus/processed_data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee01258",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc08320",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "50d1e6d1",
   "metadata": {},
   "source": [
    "## Dataset prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d603af4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5351373a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fea789fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "DATA_DIR = \"/mnt/data-ssd-1/data/supreme_corpus/processed_data/\"\n",
    "\n",
    "with open(DATA_DIR + \"segments_meta.jsonl\") as f:\n",
    "    segment_metas = [json.loads(e) for e in f]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cebc107b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b79a39e4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4458d146",
   "metadata": {},
   "source": [
    "## (optional) listen to samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702f906a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(DATA_DIR + \"segments_meta.jsonl\") as f:\n",
    "#     segment_metas = [json.loads(e) for e in f]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0185b364",
   "metadata": {},
   "outputs": [],
   "source": [
    "# idx_list = list(range(len(segment_metas)))\n",
    "\n",
    "# random.shuffle(idx_list)\n",
    "# for idx in idx_list[:5]:\n",
    "#     segment_meta = segment_metas[idx]\n",
    "#     Audio.from_file(DATA_DIR + segment_meta[\"path\"]).play()\n",
    "#     print(idx)\n",
    "#     for d in segment_meta[\"transcript\"]:\n",
    "#         print(d[\"speaker\"])\n",
    "#         print(d[\"text\"])\n",
    "#     print(\"-\"*10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd4fbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Found one slice with extra word at end: But,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66bdc29",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd1017be",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a6bd144f",
   "metadata": {},
   "source": [
    "## Add disfluencies?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2325a3c5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99a28d27",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "126acd9e",
   "metadata": {},
   "source": [
    "### Normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a94a87f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.notebook import transcribe\n",
    "from suno_utils.datasets.supreme_court.parser import _format_speaker_turns_meta\n",
    "from suno_utils.utils.parser import segments_to_tokens\n",
    "from suno_utils.utils.text import normalize_whitespace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "56e94dc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _speaker_turns_to_tokens(speaker_turns):\n",
    "    speaker_names, tokens = segments_to_tokens(_format_speaker_turns_meta([\n",
    "        (e[\"speaker\"], e[\"text\"]) for e in speaker_turns\n",
    "    ]), included_tags=[\"hesitation\", \"laughter\"])\n",
    "    for t in tokens:\n",
    "        if \"speaker_id\" in t:\n",
    "            t[\"speaker_id\"] = speaker_names.get(t[\"speaker_id\"])\n",
    "        if t[\"type\"] == \"tag\":\n",
    "            if t[\"value\"] == \"hesitation\":\n",
    "                t[\"value\"] = \"--\"\n",
    "            else:\n",
    "                t[\"value\"] = \"[{}]\".format(t[\"value\"])\n",
    "    return tokens\n",
    "\n",
    "# create token representation (-- and [laughter])\n",
    "for segment_meta in segment_metas:\n",
    "    text = \" \".join([e[\"text\"] for e in segment_meta[\"transcript\"]])\n",
    "    tokens = _speaker_turns_to_tokens(segment_meta[\"transcript\"])\n",
    "    segment_meta[\"transcript\"] = {\n",
    "        \"text\": text,\n",
    "        \"tokens\": tokens,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2ee180",
   "metadata": {},
   "outputs": [],
   "source": [
    "asr_predictions = transcribe([DATA_DIR + e[\"path\"] for e in segment_metas], n_workers=8, n_gpu=1)\n",
    "# with open(\"tmp.json\") as f:\n",
    "#     asr_predictions = json.load(f)\n",
    "\n",
    "# remove disfluencies\n",
    "asr_predictions = [\n",
    "    normalize_whitespace(re.sub(r\"\\b(um|uh)\\b\", \" \", s))\n",
    "    for s in asr_predictions\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b9628626",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.alignment import format_tokens_as_turns\n",
    "\n",
    "# remove meta tags but keep speakers as segments\n",
    "transcript_turns = []\n",
    "for segment_meta in segment_metas:\n",
    "    turns = []\n",
    "    prev_speaker_id = None\n",
    "    tmp_tokens = []\n",
    "    for t in segment_meta[\"transcript\"][\"tokens\"]:\n",
    "        if t[\"type\"] != \"text\":\n",
    "            continue\n",
    "        if t[\"speaker_id\"] != prev_speaker_id:\n",
    "            if len(tmp_tokens) > 0:\n",
    "                turns.append(\" \".join(tmp_tokens))\n",
    "                tmp_tokens = []\n",
    "            prev_speaker_id = t[\"speaker_id\"]\n",
    "        tmp_tokens.append(t[\"value\"])\n",
    "    if len(tmp_tokens) > 0:\n",
    "        turns.append(\" \".join(tmp_tokens))\n",
    "    transcript_turns.append(turns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f40cfa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text_normalizer import normalize, normalize_mp\n",
    "\n",
    "# asr/norm does mr instead of mister\n",
    "transcript_turns_norm = normalize_mp(transcript_turns, asr_predictions, n_cores=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "261988fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _norm_turns_to_tokens(turns_norm, tokens):\n",
    "    prev_speaker_id = None\n",
    "    turn_speaker_ids = []\n",
    "    for t in tokens:\n",
    "        if \"speaker_id\" in t and t[\"speaker_id\"] != prev_speaker_id:\n",
    "            turn_speaker_ids.append(t[\"speaker_id\"])\n",
    "            prev_speaker_id = t[\"speaker_id\"]\n",
    "    assert(len(turns_norm) == len(turn_speaker_ids))\n",
    "    tokens_norm = []\n",
    "    for text_norm, speaker_id in zip(turns_norm, turn_speaker_ids):\n",
    "        for token_text in text_norm.split():\n",
    "            tokens_norm.append({\n",
    "                \"type\": \"text\",\n",
    "                \"value\": token_text,\n",
    "                \"speaker_id\": speaker_id,\n",
    "            })\n",
    "    return tokens_norm\n",
    "\n",
    "for segment_meta, turns_norm, asr_pred in zip(segment_metas, transcript_turns_norm, asr_predictions):\n",
    "    segment_meta[\"transcript\"][\"text_norm\"] = \" \".join(turns_norm)\n",
    "    tokens_norm = _norm_turns_to_tokens(turns_norm, segment_meta[\"transcript\"][\"tokens\"])\n",
    "    segment_meta[\"transcript\"][\"tokens_norm\"] = tokens_norm\n",
    "    segment_meta[\"transcript\"][\"asr_transcript\"] = asr_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "0d9d19e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ~1.3Gb for supreme, save as gz?\n",
    "with open(DATA_DIR + \"segment_meta_post_norm.json\", \"w\") as f:\n",
    "    json.dump(segment_metas, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0afc53",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c69286b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f03f1240",
   "metadata": {},
   "source": [
    "## Alignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0638813b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: align timestamps next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c7fc157d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(DATA_DIR + \"segment_meta_post_norm.json\") as f:\n",
    "    segment_metas = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a6906b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.web.mfa import align_tokens, align_text\n",
    "\n",
    "conda_sh_filepath = \"/home/georg/anaconda3/etc/profile.d/conda.sh\"\n",
    "env_name = \"mfa\"\n",
    "\n",
    "audio_filepaths = [DATA_DIR + e[\"path\"] for e in segment_metas]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "22dab3e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0/100 failed.\n"
     ]
    }
   ],
   "source": [
    "transcripts_tokens = [e[\"transcript\"][\"tokens\"] for e in segment_metas]\n",
    "word_aligments = align_tokens(\n",
    "    audio_filepaths, transcripts_tokens, \n",
    "#     audio_filepaths[:100], transcripts_tokens[:100], \n",
    "    conda_sh_filepath, env_name, num_cores=40\n",
    ")\n",
    "n_failed = len([e for e in word_aligments if e is None])\n",
    "print(\"{}/{} failed.\".format(n_failed, len(word_aligments)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ed80c443",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0/100 failed.\n"
     ]
    }
   ],
   "source": [
    "# transcripts_norm = [e[\"transcript\"][\"text_norm\"] for e in segment_metas]\n",
    "# word_aligments_norm = align_text(\n",
    "#     audio_filepaths[:100], transcripts_norm[:100], \n",
    "#     conda_sh_filepath, env_name, num_cores=40\n",
    "# )\n",
    "\n",
    "transcripts_tokens_norm = [e[\"transcript\"][\"tokens_norm\"] for e in segment_metas]\n",
    "word_aligments_norm  = align_tokens(\n",
    "    audio_filepaths, transcripts_tokens_norm,\n",
    "#     audio_filepaths[:100], transcripts_tokens_norm[:100],\n",
    "    conda_sh_filepath, env_name, num_cores=40\n",
    ")\n",
    "n_failed = len([e for e in word_aligments_norm  if e is None])\n",
    "print(\"{}/{} failed.\".format(n_failed, len(word_aligments_norm)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "73243e4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# zip time info into tokens instead of adding\n",
    "for segment_meta, word_aligment, word_aligment_norm in zip(segment_metas, word_aligments, word_aligments_norm):\n",
    "    if word_aligment is None or word_aligment_norm is None:\n",
    "        segment_meta[\"ignore_segment\"] = True\n",
    "        continue\n",
    "    segment_meta[\"ignore_segment\"] = False\n",
    "    assert(len(segment_meta[\"transcript\"][\"tokens\"]) == len(word_aligment))\n",
    "    assert(len(segment_meta[\"transcript\"][\"tokens_norm\"]) == len(word_aligment_norm))\n",
    "    for t, ts in zip(segment_meta[\"transcript\"][\"tokens\"], word_aligment):\n",
    "        if ts is None:\n",
    "            assert(t[\"type\"] != \"text\")\n",
    "            continue\n",
    "        t[\"start_s\"] = ts[0]\n",
    "        t[\"end_s\"] = ts[1]\n",
    "    for t, ts in zip(segment_meta[\"transcript\"][\"tokens_norm\"], word_aligment_norm):\n",
    "        t[\"start_s\"] = ts[0]\n",
    "        t[\"end_s\"] = ts[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8cdbb489",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(DATA_DIR + \"segment_meta_post_align.json\", \"w\") as f:\n",
    "    json.dump(segment_metas, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df497d46",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cbc4d2a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b3426ad8",
   "metadata": {},
   "source": [
    "## Filter slices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b00a6b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# populate ignore_segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5537e853",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check for failures during align, maybe norm failures?\n",
    "# check for edge mistakes\n",
    "# check for high wer?\n",
    "# create nemo dataset stats pdf\n",
    "# also make sample data and speaker meta etc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9387ff11",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "094dca92",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8984f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ['source_audio_offset_s',\n",
    "#  'id',\n",
    "#  'uri',\n",
    "#  'duration_s',\n",
    "#  'source_audio_uri',\n",
    "#  'source_audio_id',\n",
    "#  'source_transcript_id',\n",
    "#  'source_transcript_uri',\n",
    "#  'transcript_normalized', # dict: text, tokens\n",
    "#  'transcript, # dict text, tokens\n",
    "#  'transcript_asr', # dict: model_id, transcipt\n",
    "# ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c328c623",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: make diariation dataset??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a928f2ca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "651a9855",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "2ccf2eb3",
   "metadata": {},
   "source": [
    "## Adjust samples data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "46cf90d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/home/georg/sample_manifest_web.json\") as f:\n",
    "    d = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "dff66557",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in d.items():\n",
    "    for n, e in enumerate(v):\n",
    "        seen_names = set()\n",
    "        speaker_names = []\n",
    "        for ee in e[\"segments\"]:\n",
    "            if ee[\"speaker\"] not in seen_names:\n",
    "                speaker_names.append(ee[\"speaker\"])\n",
    "                seen_names.add(ee[\"speaker\"])\n",
    "        speaker_rename_map = {}\n",
    "        n_justice = 1\n",
    "        n_guest = 1\n",
    "        for speaker_name in speaker_names:\n",
    "            if speaker_name is None:\n",
    "                continue\n",
    "            elif \"Justice\" in speaker_name:\n",
    "                if speaker_name not in speaker_rename_map:\n",
    "                    speaker_rename_map[speaker_name] = f\"Justice {n_justice}\"\n",
    "                    n_justice += 1\n",
    "            else:\n",
    "                if speaker_name not in speaker_rename_map:\n",
    "                    speaker_rename_map[speaker_name] = f\"Guest {n_guest}\"\n",
    "                    n_guest += 1\n",
    "        for ee in e[\"segments\"]:\n",
    "            ee[\"speaker\"] = speaker_rename_map.get(ee[\"speaker\"])\n",
    "        for ee in e[\"segments_normalized\"]:\n",
    "            ee[\"speaker\"] = speaker_rename_map.get(ee[\"speaker\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "e226875e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/home/georg/new_sample_manifest_web.json\", \"w\") as f:\n",
    "    json.dump(d, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "922386c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# s3cmd put /home/georg/new_sample_manifest_web.json s3://suno-static-public/datasets/supreme-corpus/sample-data/sample_manifest_web.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d9791e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62bd21e9",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}