{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e484d835",
   "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": 2,
   "id": "ba102c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "77b57ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.asr import transcribe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "efed8363",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('', (0.0, 1.44), 0.999),\n",
       " ('president', (1.44, 1.8), 0.915),\n",
       " (' ', (1.8, 1.96), 0.998),\n",
       " ('joe', (1.96, 2.16), 0.922),\n",
       " (' ', (2.16, 2.2), 0.947),\n",
       " ('biden', (2.2, 2.6), 0.945),\n",
       " (' ', (2.6, 2.76), 0.999),\n",
       " ('backed', (2.76, 3.04), 0.994),\n",
       " ('by', (3.04, 3.16), 0.937),\n",
       " (' ', (3.16, 3.2), 0.967),\n",
       " ('the', (3.2, 3.24), 0.997),\n",
       " (' ', (3.24, 3.28), 0.954),\n",
       " ('full', (3.28, 3.48), 0.97),\n",
       " (' ', (3.48, 3.56), 0.999),\n",
       " ('symbolic', (3.56, 4.04), 0.992),\n",
       " (' ', (4.04, 4.12), 0.996),\n",
       " ('power', (4.12, 4.36), 0.991),\n",
       " (' ', (4.36, 4.44), 0.989),\n",
       " ('of', (4.44, 4.48), 0.999),\n",
       " (' ', (4.48, 4.52), 0.997),\n",
       " ('the', (4.52, 4.56), 0.995),\n",
       " (' ', (4.56, 4.64), 0.993),\n",
       " ('western', (4.64, 5.0), 0.99),\n",
       " (' ', (5.0, 5.08), 0.996),\n",
       " ('alliance', (5.08, 5.48), 0.991),\n",
       " (' ', (5.48, 5.72), 1.0),\n",
       " ('is', (5.72, 5.76), 0.998),\n",
       " (' ', (5.76, 5.84), 0.998),\n",
       " ('locked', (5.84, 6.15), 0.99),\n",
       " ('in', (6.15, 6.19), 0.999),\n",
       " (' ', (6.19, 6.27), 0.959),\n",
       " ('a', (6.27, 6.31), 0.991),\n",
       " (' ', (6.31, 6.35), 0.995),\n",
       " ('showdown', (6.35, 6.83), 0.954),\n",
       " (' ', (6.83, 6.87), 0.999),\n",
       " ('with', (6.87, 6.95), 0.992),\n",
       " (' ', (6.95, 7.07), 0.969),\n",
       " ('russian', (7.07, 7.39), 0.998),\n",
       " (' ', (7.39, 7.43), 0.991),\n",
       " ('president', (7.43, 7.71), 0.996),\n",
       " (' ', (7.71, 7.79), 0.736),\n",
       " ('vladimir', (7.79, 8.31), 0.927),\n",
       " (' ', (8.31, 8.35), 0.994),\n",
       " ('putin', (8.35, 8.75), 0.978),\n",
       " (' ', (8.75, 8.87), 0.999),\n",
       " ('who', (8.87, 9.07), 0.97),\n",
       " (' ', (9.07, 9.11), 0.998),\n",
       " ('is', (9.11, 9.15), 0.997),\n",
       " (' ', (9.15, 9.23), 0.999),\n",
       " ('using', (9.23, 9.51), 0.972),\n",
       " (' ', (9.51, 9.63), 0.998),\n",
       " ('ukraine', (9.63, 10.11), 0.944),\n",
       " (' ', (10.11, 10.27), 1.0),\n",
       " ('as', (10.27, 10.39), 0.996),\n",
       " ('a', (10.39, 10.47), 0.962),\n",
       " (' ', (10.47, 10.55), 0.984),\n",
       " ('hostage', (10.55, 11.11), 0.992),\n",
       " (' ', (11.11, 11.23), 0.999),\n",
       " ('to', (11.23, 11.31), 0.878),\n",
       " (' ', (11.31, 11.35), 1.0),\n",
       " ('try', (11.35, 11.55), 0.995),\n",
       " (' ', (11.55, 11.59), 0.999),\n",
       " ('to', (11.59, 11.67), 0.938),\n",
       " (' ', (11.67, 11.75), 1.0),\n",
       " ('force', (11.75, 11.95), 0.989),\n",
       " (' ', (11.95, 12.03), 1.0),\n",
       " ('the', (12.03, 12.11), 0.992),\n",
       " (' ', (12.11, 12.19), 1.0),\n",
       " ('us', (12.19, 12.35), 0.789),\n",
       " (' ', (12.35, 12.51), 0.996),\n",
       " ('to', (12.51, 12.55), 0.999),\n",
       " (' ', (12.55, 12.67), 0.943),\n",
       " ('renegotiate', (12.67, 13.51), 0.997),\n",
       " (' ', (13.51, 13.67), 0.999),\n",
       " ('the', (13.67, 13.71), 0.999),\n",
       " (' ', (13.71, 13.79), 0.992),\n",
       " ('settled', (13.79, 14.19), 0.947),\n",
       " (' ', (14.19, 14.23), 0.995),\n",
       " ('outcome', (14.23, 14.63), 0.942),\n",
       " (' ', (14.63, 14.71), 0.999),\n",
       " ('of', (14.71, 14.79), 0.996),\n",
       " (' ', (14.79, 14.83), 0.998),\n",
       " ('the', (14.83, 14.91), 0.959),\n",
       " (' ', (14.91, 14.95), 1.0),\n",
       " ('cold', (14.95, 15.23), 0.974),\n",
       " (' ', (15.23, 15.31), 1.0),\n",
       " ('war', (15.31, 15.47), 0.999),\n",
       " ('', (15.47, 16.07), 1.0)]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transcribe(\"/home/georg/data/sample_audio/russia.wav\", mode=\"tokens\", use_gpu=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f7f1d3d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d72c7c5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c2c95a7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df27790d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ed625beb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No GPU found even though requested, using CPU.\n"
     ]
    }
   ],
   "source": [
    "vocab, logits = transcribe(\"/home/georg/data/sample_audio/russia.wav\", mode=\"logits\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "106a64c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyctcdecode import build_ctcdecoder\n",
    "decoder = build_ctcdecoder(vocab)\n",
    "_, _, tokens_with_frames, _, _ = decoder.decode_beams(logits, beam_width=2)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "a81f8e26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(402, 129)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "33e70495",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('', (0, 36)),\n",
       " ('president', (36, 45)),\n",
       " (' ', (45, 49)),\n",
       " ('joe', (49, 54)),\n",
       " (' ', (54, 55)),\n",
       " ('biden', (55, 65)),\n",
       " (' ', (65, 69)),\n",
       " ('backed', (69, 76))]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cur_frame = 0\n",
    "tokens_with_frames_and_space = []\n",
    "for token, (start_f, end_f) in tokens_with_frames:\n",
    "    if start_f - cur_frame > 0:\n",
    "        space_token = \"\" if cur_frame == 0 else \" \"\n",
    "        tokens_with_frames_and_space.append((space_token, (cur_frame, start_f)))\n",
    "    tokens_with_frames_and_space.append((token, (start_f, end_f)))\n",
    "    cur_frame = end_f\n",
    "if len(logits) > cur_frame:\n",
    "    tokens_with_frames_and_space.append((\"\", (cur_frame, len(logits))))\n",
    "tokens_with_frames_and_space[:8]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "619dc54b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(402, 129)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "7a2b843c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.special import softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "eca7556c",
   "metadata": {},
   "outputs": [],
   "source": [
    "probs = softmax(logits, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5aee6ba2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9988637"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probs[0:3].max(axis=1).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6918e05b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cef3e7d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4f8d5b8",
   "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
}