02_Preprocess_02_prepare_fa_dialogue_auto.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5/10 (Fri) | Preprocess\n",
    "\n",
    "# Forced Alignment Preparation for Dialogue Data (ASR Transcript)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Introduction\n",
    "\n",
    "Thie notebook generates dialogue transcripts and split full audio data for wav2vec forced alignment.\n",
    "The preparation process consists of the following stages.\n",
    "\n",
    "1. Get user ids\n",
    "2. Load an audio file and a transcript text corresponding to a user id\n",
    "3. Get a turn-level ids, transcript, and start and end time iteratively\n",
    "4. Convert transcript to lower-case string with star tokens\n",
    "5. Split audio speech by the start and end time\n",
    "6. Save the turn-level transcript and split audio\n",
    "\n",
    "Before starting the process, the following code block loads required packages and define gloabl variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple, Generator\n",
    "from pathlib import Path\n",
    "import re\n",
    "\n",
    "import pandas as pd\n",
    "from pydub import AudioSegment\n",
    "\n",
    "TRANSCRIPT_DIR = Path(\"/home/matsuura/Development/app/feature_extraction_api/experiment/data/WoZ_Interview/02_Rev_Transcript\")\n",
    "AUDIO_DIR = Path(\"/home/matsuura/Development/Datasets/teai-woz-2021-03/wav\")\n",
    "SAVE_DIR = Path(\"/home/matsuura/Development/app/feature_extraction_api/experiment/data/WoZ_Interview/03_FA_Audio_Transcript_Auto\")\n",
    "\n",
    "PUNCTUATIONS = [\".\", \",\", \":\", \"?\", \"!\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 2. Define Functions\n",
    "\n",
    "This section defines functions to complete the preparation.\n",
    "The following code block defines a generator of user ids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def user_id_generator() -> Generator[str, None, None]:\n",
    "    for uid in range(1, 86):\n",
    "        uid = str(uid).zfill(3)\n",
    "\n",
    "        yield uid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code block defines functions to load a transcript and an audio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_transcript(uid: str) -> pd.DataFrame:\n",
    "    csv_path = TRANSCRIPT_DIR / f\"{uid}.csv\"\n",
    "    df_transcript = pd.read_csv(csv_path, index_col=0)\n",
    "\n",
    "    return df_transcript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_audio(uid: str) -> AudioSegment:\n",
    "    if int(uid) == 19:\n",
    "        uid = \"086\"\n",
    "\n",
    "    if int(uid) < 56:\n",
    "        audio_path = AUDIO_DIR / f\"{uid}-user.mp4\"\n",
    "    else:\n",
    "        audio_path = AUDIO_DIR / f\"{uid}-user.m4a\"\n",
    "\n",
    "    audio = AudioSegment.from_file(audio_path)\n",
    "\n",
    "    return audio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code block defines a generator of turn-level ids, transcripts, and start and end times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def turn_level_info_generator(\n",
    "        df_transcript: pd.DataFrame\n",
    ") -> Generator[Tuple[str, str, int, int], None, None]:\n",
    "    user_mask = (df_transcript[\"speaker\"] == \"user\")\n",
    "    \n",
    "    intro_mask = (df_transcript[\"topic\"] == \"intro\")\n",
    "    closing_mask = (df_transcript[\"topic\"] == \"closing\")\n",
    "    topic_mask = intro_mask | closing_mask\n",
    "\n",
    "    mask = user_mask & (~topic_mask)\n",
    "\n",
    "    df_transcript_masked = df_transcript[mask]\n",
    "    \n",
    "    for idx in df_transcript_masked.index:\n",
    "        transcript = df_transcript_masked.at[idx, \"transcript\"]\n",
    "        start_time = df_transcript_masked.at[idx, \"start_time\"]\n",
    "        end_time = df_transcript_masked.at[idx, \"end_time\"]\n",
    "\n",
    "        turn_id = str(idx).zfill(3)\n",
    "\n",
    "        yield turn_id, transcript, start_time, end_time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code block defines a function to convert original transcripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_fa_transcript(transcript: str) -> str:\n",
    "    # 1. change an inaudible tags to star tokens\n",
    "    fa_transcript = transcript.replace(\"<inaudible>\", \"*\")\n",
    "\n",
    "    # 2. remove punctuations\n",
    "    for punct in PUNCTUATIONS:\n",
    "       fa_transcript = fa_transcript.replace(punct, \"\") \n",
    "    fa_transcript = fa_transcript.replace(\"-\", \"\")\n",
    "\n",
    "    # 3. change numbers to star tokens\n",
    "    number_pattern = r\"\\b\\w*\\d\\w*\\b\"\n",
    "    fa_transcript = re.sub(number_pattern, \"*\", fa_transcript)\n",
    "\n",
    "    # 4. remove other tags\n",
    "    tag_pattern = r\"\\<.*?\\>\"\n",
    "    fa_transcript = re.sub(tag_pattern, \" \", fa_transcript) \n",
    "\n",
    "    # 5. lower transcript\n",
    "    fa_transcript = fa_transcript.lower()\n",
    "\n",
    "    # 6. remove extra pauses\n",
    "    while \"  \" in fa_transcript:\n",
    "        fa_transcript = fa_transcript.replace(\"  \", \" \")\n",
    "\n",
    "    if fa_transcript[0] == \" \":\n",
    "        fa_transcript = fa_transcript[1:]\n",
    "    if fa_transcript[-1] == \" \":\n",
    "        fa_transcript = fa_transcript[:-1]\n",
    "\n",
    "    return fa_transcript"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code block defines a function to split AudioSegment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_audio(audio: AudioSegment, start_time: int, end_time: int) -> AudioSegment:\n",
    "    return audio[start_time:end_time]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code block defines a function to save transcript and audio files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_transcript(uid: str, turn_id: str, fa_transcript: str) -> None:\n",
    "    save_path = SAVE_DIR / f\"{uid}_{turn_id}.txt\"\n",
    "    with open(save_path, \"w\") as f:\n",
    "        f.write(fa_transcript)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_audio(uid: str, turn_id: str, audio: AudioSegment) -> None:\n",
    "    save_path = SAVE_DIR / f\"{uid}_{turn_id}.wav\"\n",
    "    audio.export(save_path, format=\"wav\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 3. Conduct Forced Alignment Preparation\n",
    "\n",
    "the following code block conduct the preparation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for uid in user_id_generator():\n",
    "    df_transcript = load_transcript(uid)\n",
    "    audio = load_audio(uid)\n",
    "\n",
    "    for tid, transcript, start_time, end_time in turn_level_info_generator(df_transcript):\n",
    "        if isinstance(transcript, float):\n",
    "            continue # skip empty transcripts\n",
    "\n",
    "        fa_transcript = convert_fa_transcript(transcript)\n",
    "        fa_audio = split_audio(audio, start_time, end_time)\n",
    "\n",
    "        save_transcript(uid, tid, fa_transcript)\n",
    "        save_audio(uid, tid, fa_audio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "teai-incremental-classifier",
   "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.9.undefined"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}