prepare_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install transformers datasets \n",
    "! pip install torchaudio\n",
    "! pip install librosa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "print(torch.cuda.is_available())  # Should return True if GPUs are available\n",
    "print(torch.cuda.device_count())  # Number of available GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load the dataset\n",
    "dataset = load_dataset(\"ganga4364/Dilgo-Khyentse-Rinpoche-dataset\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Trainer, TrainingArguments\n",
    "\n",
    "# Load pre-trained Wav2Vec2 model and processor\n",
    "model_name = \"ganga4364/mms_300_v4.96000\"\n",
    "processor = Wav2Vec2Processor.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchaudio.transforms import Resample\n",
    "import torchaudio\n",
    "\n",
    "def speech_file_to_array_fn(batch):\n",
    "    # print(batch)\n",
    "    speech_array, sampling_rate = torchaudio.load(batch[\"path\"])\n",
    "    # print(speech_array.shape, sampling_rate)\n",
    "    if sampling_rate != 16000:\n",
    "        print(\"resampling\")\n",
    "        resampler = Resample(orig_freq=sampling_rate, new_freq=16000)\n",
    "        speech_array = resampler(speech_array)\n",
    "        sampling_rate = 16000\n",
    "    \n",
    "    # print(speech_array.shape, sampling_rate)\n",
    "    batch[\"speech\"] = speech_array[0].numpy()\n",
    "    batch[\"sampling_rate\"] = sampling_rate\n",
    "    batch[\"target_text\"] = batch[\"uni\"]\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def prepare_dataset(batch):\n",
    "    batch[\"input_values\"] = processor(batch[\"speech\"], sampling_rate=batch[\"sampling_rate\"]).input_values\n",
    "    # reshape to (n,)\n",
    "    batch[\"input_values\"] = np.squeeze(batch[\"input_values\"])\n",
    "    # if batch[\"sampling_rate\"] != 16000:\n",
    "    #     print(\"sampling rate not 16k\", batch)\n",
    "    \n",
    "    # with processor.as_target_processor():\n",
    "    #     batch[\"labels\"] = processor(batch[\"target_text\"]).input_ids\n",
    "\n",
    "    batch[\"labels\"] = processor(text=batch[\"target_text\"]).input_ids\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = dataset['train']\n",
    "dataset_val = dataset['validation']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assuming dataset_test and dataset_val are datasets.Dataset objects\n",
    "dataset_val = dataset_val.map(lambda x: {'path': f\"/data/volume/wav_16k/{x['file_name']}.wav\"})\n",
    "dataset_train = dataset_train.map(lambda x: {'path': f\"/data/volume/wav_16k/{x['file_name']}.wav\"})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "\n",
    "df = pd.read_csv('04_bad_apples.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Get the list of filenames to remove\n",
    "filenames_to_remove = df['file_name'].tolist()  # Your dataframe containing the filenames to remove\n",
    "\n",
    "# Step 2: Define a function to check if a file_name starts with any of the filenames to remove\n",
    "def should_keep(example):\n",
    "    return not any(example['file_name'].startswith(f) for f in filenames_to_remove)\n",
    "\n",
    "# Step 3: Use the map function to filter the dataset\n",
    "dataset_train = dataset_train.filter(should_keep)\n",
    "dataset_val = dataset_val.filter(should_keep)\n",
    "\n",
    "# Optional: Save the filtered dataset\n",
    "# dataset_train_filtered.save_to_disk('filtered_dataset')  # Uncomment to save it to disk\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = dataset_train.map(speech_file_to_array_fn, remove_columns=dataset_train.column_names)\n",
    "dataset_train = dataset_train.map(prepare_dataset, remove_columns=dataset_train.column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train.save_to_disk('/data/volume/wav2vec2/train_prepare_dataset.arrow')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use num_proc=3 to enable multiprocessing\n",
    "dataset_val = dataset_val.map(speech_file_to_array_fn, remove_columns=dataset_val.column_names, num_proc=5)\n",
    "dataset_val = dataset_val.map(prepare_dataset, remove_columns=dataset_val.column_names, num_proc=5)\n",
    "# Save the training, validation, and test datasets to disk\n",
    "dataset_val.save_to_disk('/data/volume/wav2vec2/val_prepare_dataset.arrow')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupyter-base",
   "language": "python",
   "name": "jupyter-base"
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}