{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42fa1670",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: fix proper indicator token stripping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "abb40508",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from suno_utils.web.rev import parse_transcript\n",
    "from suno_utils.customers.sanas.pipeline import get_segments\n",
    "\n",
    "REV_DIR = '/mnt/data-ssd-1/data/academia/hf_paper/rev'\n",
    "with open(os.path.join(REV_DIR, f\"concat_meta.json\")) as f:\n",
    "    rev_audio_meta = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3e5bf060",
   "metadata": {},
   "outputs": [],
   "source": [
    "transcripts_container = {}\n",
    "for group_tag, segments_meta in rev_audio_meta.items():\n",
    "    if group_tag != \"ami_clean\":\n",
    "        continue\n",
    "    transcript_filepath = os.path.join(REV_DIR, \"transcript\", f\"{group_tag}.json\")\n",
    "    with open(transcript_filepath) as f:\n",
    "        raw_transcript = json.load(f)\n",
    "    transcript = parse_transcript(\n",
    "        raw_transcript, anonymize_speakers=True\n",
    "    )\n",
    "    n_segments = segments_meta[-1][\"segment_number\"] + 1\n",
    "    indexed_token_segments = get_segments(transcript, n_segments)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "64354ebc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'00. Yeah.'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indexed_token_segments[0][1].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "628a3722",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "e9dba507",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Counter(indicator_segments)\n",
    "c.most_common(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "11ada9f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Speaker 12', 440), ('Speaker 0', 385)]"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.most_common(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "ad0ac94f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "unreliable global indicator speaker\n"
     ]
    }
   ],
   "source": [
    "        c = Counter(indicator_segments)\n",
    "        sl = c.most_common(2)\n",
    "        if len(sl) == 2 and sl[0][1] * 0.8 < sl[1][1]:\n",
    "            logger.warning(\"unreliable global indicator speaker\")\n",
    "        indicator_speaker_id = sl[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "0de74ae6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Speaker 12'"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indicator_speaker_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "675be8f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "440"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sl[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "62a85d9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "385"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sl[1][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14703dcc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "517833d5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c344429",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1cbebbbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.customers.sanas.pipeline import *\n",
    "from suno_utils.customers.sanas.pipeline import _is_numeric, _digitify, _find_valid_indicators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f2a56041",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "return cleaned_token_segments\n"
     ]
    }
   ],
   "source": [
    "rev_transcript = transcript\n",
    "global_indicator_speaker = False\n",
    "\n",
    "if global_indicator_speaker:\n",
    "    # find indicator speaker\n",
    "    indicator_segments = [\n",
    "        t[\"speaker_id\"]\n",
    "        for t in rev_transcript.speaker_turns\n",
    "        if _is_numeric(t[\"plaintext\"])\n",
    "    ]\n",
    "    if len(indicator_segments) == 0:\n",
    "        logger.warning(\"no indicators found\")\n",
    "        print(\"return []\")\n",
    "    indicator_speaker_id = mode(indicator_segments)\n",
    "else:\n",
    "    indicator_speaker_id = None\n",
    "# parse numbers for indicator speaker\n",
    "indicator_tokens = []\n",
    "prev_n = -2\n",
    "last_speaker_id = None\n",
    "for n, token in enumerate(rev_transcript.tokens):\n",
    "    # if we know indicator speaker id then skip others\n",
    "    if indicator_speaker_id is not None and indicator_speaker_id != token.speaker_id:\n",
    "        last_speaker_id = token.speaker_id\n",
    "        continue\n",
    "    parsed_number = _digitify(token.value)\n",
    "    if len(parsed_number) == 0:\n",
    "        last_speaker_id = token.speaker_id\n",
    "        continue\n",
    "    # if multiple tokens in a row pass then add to previous\n",
    "    if n - prev_n == 1 and last_speaker_id == token.speaker_id:\n",
    "        indicator_tokens[-1] = (\n",
    "            indicator_tokens[-1][0],\n",
    "            indicator_tokens[-1][1] + parsed_number,\n",
    "        )\n",
    "    elif last_speaker_id != token.speaker_id:\n",
    "        indicator_tokens.append((n, parsed_number))\n",
    "    else:\n",
    "        # TODO: if not new speaker then we don't add?\n",
    "        last_speaker_id = token.speaker_id\n",
    "        continue\n",
    "    prev_n = n\n",
    "    last_speaker_id = token.speaker_id\n",
    "\n",
    "is_indicators_valid = _find_valid_indicators(\n",
    "    [s for _, s in indicator_tokens],\n",
    "    [str(n).zfill(2) for n in range(n_segments)]\n",
    ")\n",
    "correct_indicators = [\n",
    "    (int(s), idx) for (idx, s), is_valid in zip(indicator_tokens, is_indicators_valid)\n",
    "    if is_valid\n",
    "]\n",
    "\n",
    "# fix missing indicators if possible\n",
    "missing_indicators = sorted(list(\n",
    "    set(range(n_segments)) - set([n for n, _ in correct_indicators])\n",
    "))\n",
    "added_indicators = []\n",
    "for missing_indicator in missing_indicators:\n",
    "    # obtain search index\n",
    "    l = [e for e in correct_indicators if e[0] < missing_indicator]\n",
    "    if len(l) > 0:\n",
    "        start_idx = l[-1][-1]\n",
    "    else:\n",
    "        start_idx = 0\n",
    "    l = [e for e in correct_indicators if e[0] > missing_indicator]\n",
    "    if len(l) > 0:\n",
    "        end_idx = l[0][-1]\n",
    "    else:\n",
    "        end_idx = len(rev_transcript.tokens)\n",
    "    # find index within\n",
    "    missing_indicator_str = str(missing_indicator).zfill(2)\n",
    "    indicator_token_idx = None\n",
    "    n_matched = 0\n",
    "    for n, t in enumerate(rev_transcript.tokens[start_idx:end_idx]):\n",
    "        if _digitify(t.value) == missing_indicator_str[n_matched]:\n",
    "            if n_matched == 0:\n",
    "                indicator_token_idx = start_idx + n\n",
    "            n_matched += 1\n",
    "        elif len(_digitify(t.value)) > 0 and int(_digitify(t.value)) == missing_indicator:\n",
    "            if n_matched == 0:\n",
    "                indicator_token_idx = start_idx + n\n",
    "            n_matched = len(missing_indicator_str)\n",
    "        else:\n",
    "            n_matched = 0\n",
    "            indicator_token_idx = None\n",
    "        if n_matched == len(missing_indicator_str):\n",
    "            break\n",
    "    if indicator_token_idx is not None:\n",
    "        added_indicators.append((int(missing_indicator_str), indicator_token_idx))\n",
    "correct_indicators = sorted(correct_indicators + added_indicators, key=lambda x: x[0])\n",
    "\n",
    "# get relevant tokens for all consecutive ones\n",
    "token_segments = []\n",
    "for (indicator_index_l, token_index_l), (indicator_index_r, token_index_r) in zip(\n",
    "    correct_indicators[:-1], correct_indicators[1:]\n",
    "):\n",
    "    if indicator_index_l == indicator_index_r - 1:\n",
    "        relevant_tokens = [\n",
    "            t\n",
    "            for t in rev_transcript.tokens[token_index_l:token_index_r]\n",
    "            if indicator_speaker_id is None or indicator_speaker_id != t.speaker_id\n",
    "        ]\n",
    "        token_segments.append((indicator_index_l, Tokens(relevant_tokens)))\n",
    "# check last one\n",
    "indicator_index_l, token_index_l = correct_indicators[-1]\n",
    "if indicator_index_l == n_segments - 1:\n",
    "    relevant_tokens = [\n",
    "        t\n",
    "        for t in rev_transcript.tokens[token_index_l:]\n",
    "        if indicator_speaker_id is None or indicator_speaker_id != t.speaker_id\n",
    "    ]\n",
    "    token_segments.append((indicator_index_l, Tokens(relevant_tokens)))\n",
    "# if we don't have a global indicator we want to remove the first speaker from each segment\n",
    "if indicator_speaker_id is None:\n",
    "    cleaned_token_segments = []\n",
    "    for indicator_idx, tokens in token_segments:\n",
    "        if len(tokens) == 0:\n",
    "            raise ValueError(\"no tokens found in indicator segment, this should not happen.\")\n",
    "        cleaned_tokens = []\n",
    "        for n, t in enumerate(tokens[1:]):\n",
    "            if t.speaker_id != tokens[0].speaker_id or len(_digitify(t.value)) == 0:\n",
    "                cleaned_tokens = tokens[n+1:]\n",
    "                break\n",
    "        cleaned_token_segments.append((indicator_idx, Tokens(cleaned_tokens)))\n",
    "else:\n",
    "    cleaned_token_segments = token_segments\n",
    "# TODO: sometimes leading '03.' is left intact. maybe wrong speaker?\n",
    "print(\"return cleaned_token_segments\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c412670",
   "metadata": {},
   "outputs": [],
   "source": [
    "    # find indicator speaker\n",
    "    indicator_segments = [\n",
    "        t[\"speaker_id\"]\n",
    "        for t in rev_transcript.speaker_turns\n",
    "        if _is_numeric(t[\"plaintext\"])\n",
    "    ]\n",
    "    if len(indicator_segments) == 0:\n",
    "        logger.warning(\"no indicators found\")\n",
    "        print(\"return []\")\n",
    "    indicator_speaker_id = mode(indicator_segments)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17d24148",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fdf787af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 0)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "correct_indicators[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "a899943f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(837, 7707)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "correct_indicators[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f1934e6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Token(value='00.', type='text', speaker_id='Speaker 0', start_s=None, end_s=None, metadata=None),\n",
       " Token(value='Yeah.', type='text', speaker_id='Speaker 1', start_s=None, end_s=None, metadata=None),\n",
       " Token(value='01.', type='text', speaker_id='Speaker 0', start_s=None, end_s=None, metadata=None),\n",
       " Token(value='So', type='text', speaker_id='Speaker 1', start_s=None, end_s=None, metadata=None),\n",
       " Token(value='I', type='text', speaker_id='Speaker 1', start_s=None, end_s=None, metadata=None)]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rev_transcript.tokens[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d235224a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Token(value='837.', type='text', speaker_id='Speaker 12', start_s=None, end_s=None, metadata=None),\n",
       " Token(value='Fine.', type='text', speaker_id='Speaker 6', start_s=None, end_s=None, metadata=None)]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rev_transcript.tokens[7707:7707+5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c590cf5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "582c99e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, Tokens(text=`00. Yeah.`))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "token_segments[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3caf5260",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(837, Tokens(text=`Fine.`))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "token_segments[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "377604e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, Tokens(text=`00. Yeah.`))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cleaned_token_segments[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "846ada11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(837, Tokens(text=`Fine.`))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cleaned_token_segments[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b96468e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, Tokens(text=`Yeah.`))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cleaned_token_segments[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "df340ab1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(837, Tokens(text=`Fine.`))"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cleaned_token_segments[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dab4db65",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d45913",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0964b67",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f79b77b9",
   "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
}