00_core.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# core\n",
    "\n",
    "> Transcribe multilingual audio files with multiple speakers and get a Pandas dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'whisper'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Input \u001b[0;32mIn [3]\u001b[0m, in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#| export\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[0;32m----> 3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mwhisper\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpydub\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AudioSegment\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpyannote\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maudio\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Pipeline\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'whisper'"
     ]
    }
   ],
   "source": [
    "#| export\n",
    "import pandas as pd\n",
    "import whisper\n",
    "from pydub import AudioSegment\n",
    "from pyannote.audio import Pipeline\n",
    "diarize_pipeline = Pipeline.from_pretrained('pyannote/speaker-diarization')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class Dictation:\n",
    "    \n",
    "    def __init__(self): pass\n",
    "    \n",
    "    def transcribe(self, file, model=\"medium\", task=\"transcribe\"):\n",
    "\n",
    "        audio = AudioSegment.from_mp3(file)\n",
    "        audio.export('audio.wav', format='wav')\n",
    "        \n",
    "        transcriber = whisper.load_model(model)\n",
    "        \n",
    "        diarization = diarize_pipeline('audio.wav')\n",
    "        \n",
    "        df = pd.DataFrame()\n",
    "        \n",
    "        count = 0\n",
    "        for turn, track, speaker in diarization.itertracks(yield_label=True):\n",
    "            df.loc[count,'start'] = turn.start\n",
    "            df.loc[count,'end'] = turn.end\n",
    "            df.loc[count,'speaker'] = speaker\n",
    "            \n",
    "            count += 1\n",
    "        \n",
    "        df['difference'] = df['end'] - df['start']\n",
    "        \n",
    "        #make a DF that merges tracks of the same speaker and tries to weed out overlaps\n",
    "        simplified_df = pd.DataFrame()\n",
    "        simplified_df.loc[0,'start'] = df.loc[0,'start']\n",
    "        simplified_df.loc[0,'speaker'] = df.loc[0,'speaker']\n",
    "\n",
    "        count = 0\n",
    "        for index, row in df.iterrows():\n",
    "            if df.loc[index,'speaker'] == simplified_df.loc[count,'speaker']:\n",
    "                last_end_time = df.loc[index,'end']\n",
    "            else:\n",
    "                #if there has been a change in speaker, check if there is an overlap.\n",
    "                #ignore overlaps if the track is less than 2s\n",
    "                #if it is longer than 2 seconds, start it where the overlap ends\n",
    "                if df.loc[index,'start'] < last_end_time:\n",
    "                    if (df.loc[index,'difference'] > 2.0) and (df.loc[index,'end'] > last_end_time):\n",
    "                        simplified_df.loc[count,'end'] = last_end_time\n",
    "                        count += 1\n",
    "                        simplified_df.loc[count,'start'] = last_end_time\n",
    "                        simplified_df.loc[count,'speaker'] = df.loc[index,'speaker']\n",
    "                        last_end_time = df.loc[index,'end']\n",
    "                    else:\n",
    "                        pass\n",
    "                else:\n",
    "                    simplified_df.loc[count,'end'] = last_end_time\n",
    "                    count += 1\n",
    "                    simplified_df.loc[count,'start'] = df.loc[index,'start']\n",
    "                    simplified_df.loc[count,'speaker'] = df.loc[index,'speaker']\n",
    "                    last_end_time = df.loc[index,'end']\n",
    "\n",
    "        simplified_df.loc[count,'end'] = df.loc[index,'end']\n",
    "        \n",
    "        df = simplified_df\n",
    "\n",
    "        #transcribe\n",
    "        for index,row in df.iterrows():\n",
    "            start = df.loc[index,'start']\n",
    "            end = df.loc[index,'end']\n",
    "            clip = audio[(start * 1000):(end * 1000)]\n",
    "            clip.export('for_transcription.wav', format='wav')\n",
    "            result = transcriber.transcribe('for_transcription.wav', task=task)\n",
    "            df.loc[index,'text'] = result['text']\n",
    "            df.loc[index,'language'] = result['language']\n",
    "        \n",
    "        self.results = df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}