src/transformers/notebooks/asr_pipelines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ASR with OpenAI/Whisper-large-v3\n",
    "\n",
    "## Build an Automatic Speech Recognition Pipeline\n",
    "\n",
    "- [Whisper (Large) on Hugging Face](https://huggingface.co/openai/whisper-large-v3 \"official openai repository\")\n",
    "- [Hugging Face course on ASR with Whisper (Small)](https://huggingface.co/learn/audio-course/chapter5/asr_models#graduation-to-seq2seq)\n",
    "\n",
    "\n",
    "#### Usage:\n",
    "Whisper large-v3 is supported in Hugging Face 🤗 Transformers through the main branch in the Transformers repo. To run the model, first install the Transformers library through the GitHub repo. For this example, we'll also install 🤗 Datasets to load toy audio dataset from the Hugging Face Hub:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that pip is the latest version:\n",
    "%pip install -U pip\n",
    "\n",
    "# Install transformers via source along with other necessary requirements for our notebook\n",
    "%pip install -U git+https://github.com/huggingface/transformers.git accelerate datasets[audio] torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model can be used with the [pipeline](https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.AutomaticSpeechRecognitionPipeline) class to transcribe audio files of arbitrary length. Transformers uses a chunked algorithm to transcribe long-form audio files, which in-practice is 9x faster than the sequential algorithm proposed by OpenAI (see Table 7 of the [Distil-Whisper paper](https://arxiv.org/abs/2311.00430)). The batch size should be set based on the specifications of your device:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import torch\n",
    "from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline\n",
    "\n",
    "# Check if a GPU is available and set the device accordingly\n",
    "# If a GPU is available, use it (cuda:0); otherwise, use the CPU\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Set the data type for tensors based on the availability of a GPU\n",
    "torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model ID for the pre-trained model\n",
    "model_id = \"openai/whisper-large-v3\"\n",
    "\n",
    "# Load the pre-trained model with specific configurations\n",
    "model = AutoModelForSpeechSeq2Seq.from_pretrained(\n",
    "    model_id, \n",
    "    torch_dtype=torch_dtype,  # Use the previously set data type for tensors\n",
    "    low_cpu_mem_usage=True,  # Optimize memory usage for CPU\n",
    "    use_safetensors=True     # Enable SafeTensors for memory optimization\n",
    ")\n",
    "\n",
    "# Move the model to the specified device (GPU or CPU)\n",
    "model.to(device)\n",
    "\n",
    "# Load the processor for the model\n",
    "processor = AutoProcessor.from_pretrained(model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up a pipeline for automatic speech recognition\n",
    "pipe = pipeline(\n",
    "    \"automatic-speech-recognition\",\n",
    "    model=model,  # Use the loaded model\n",
    "    tokenizer=processor.tokenizer,  # Use the tokenizer from the processor\n",
    "    feature_extractor=processor.feature_extractor,  # Use the feature extractor from the processor\n",
    "    max_new_tokens=128,\n",
    "    chunk_length_s=30,  # Set the chunk length for processing\n",
    "    batch_size=16,  # Set batch size\n",
    "    return_timestamps=True,  # Return timestamps for the transcriptions\n",
    "    torch_dtype=torch_dtype,  # Use the specified data type for tensors\n",
    "    device=device  # Specify the device (GPU or CPU)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Load a dataset for validation\n",
    "dataset = load_dataset(\"distil-whisper/librispeech_long\", \"clean\", split=\"validation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = dataset[0][\"audio\"]  # Get the first sample from the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the pipeline on the sample and get the result\n",
    "result = pipe(sample)\n",
    "\n",
    "# Print the recognized text from the audio\n",
    "print(result[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To transcribe a local audio file, simply pass the path to your audio file when you call the pipeline:\n",
    "\n",
    "```python\n",
    "- result = pipe(sample)\n",
    "+ result = pipe(\"audio.mp3\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Configure Target Language\n",
    "Whisper predicts the language of the source audio automatically. If the source audio language is known *a-priori*, it can be passed as an argument to the pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = pipe(sample, generate_kwargs={\"language\": \"english\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Configure Task\n",
    "By default, Whisper performs the task of speech transcription, where the source audio language is the same as the target text language. To perform speech translation, where the target text is in English, set the task to \"translate\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = pipe(sample, generate_kwargs={\"task\": \"translate\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Timestamps\n",
    "Finally, the model can be made to predict timestamps.\n",
    "\n",
    "For ***sentence-level timestamps***, pass the `return_timestamps` argument with `True`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = pipe(sample, return_timestamps=True)\n",
    "print(result[\"chunks\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***And for word-level timestamps, pass `\"Word\"`***:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = pipe(sample, return_timestamps=\"word\")\n",
    "print(result[\"chunks\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above arguments can be used in ***isolation*** or in ***combination***. \n",
    "\n",
    "For example, to perform the task of speech transcription where the source audio is in French, and we want to return sentence-level timestamps, the following can be used:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = pipe(sample, return_timestamps=True, generate_kwargs={\"language\": \"french\", \"task\": \"translate\"})\n",
    "print(result[\"chunks\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional Speed & Memory Improvements\n",
    "\n",
    "#### Flash Attention\n",
    "We recommend using [Flash-Attention 2](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#flashattention-2) if your GPU allows for it. To do so, you first need to install [Flash Attention](https://github.com/Dao-AILab/flash-attention):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install flash-attn -no-build-isolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and then all you have to do is to pass `use_flash_attention_2=True` to *`from_pretrained`*:\n",
    "\n",
    "```python\n",
    "- model = AutoModelForSpeechSeq2Seq.from_pretrained(\n",
    "    model_id, torch_dtype=torch_dtype, \n",
    "    low_cpu_mem_usage=True, use_safetensors=True)\n",
    "\n",
    "+ model = AutoModelForSpeechSeq2Seq.from_pretrained(\n",
    "    model_id, torch_dtype=torch_dtype, \n",
    "    low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=True) # Use Flash Attention 2\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a demo with Gradio\n",
    "\n",
    "We’ll define a function that takes the filepath for an audio input and passes it through the pipeline. \n",
    "\n",
    "Here, the pipeline automatically takes care of loading the audio file, resampling it to the correct sampling rate, and running inference with the model. We can then simply return the transcribed text as the output of the function. \n",
    "\n",
    "To ensure our model can handle audio inputs of arbitrary length, we’ll enable chunking:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transcribe_speech(filepath):\n",
    "    result = pipe(\n",
    "        filepath,\n",
    "        max_new_tokens=128,\n",
    "        generate_kwargs={\n",
    "            \"task\": \"transcribe\",\n",
    "            \"language\": \"english\",\n",
    "        },\n",
    "        chunk_length_s=30,\n",
    "        batch_size=8,\n",
    "    )\n",
    "    return result[\"text\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We’ll use the Gradio blocks feature to launch two tabs on our demo: one for microphone transcription, and the other for file upload."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "demo = gr.Blocks()\n",
    "\n",
    "mic_transcribe = gr.Interface(\n",
    "    fn=transcribe_speech,\n",
    "    inputs=gr.Audio(sources=\"microphone\", type=\"filepath\"),\n",
    "    outputs=gr.outputs.Textbox(),\n",
    ")\n",
    "\n",
    "file_transcribe = gr.Interface(\n",
    "    fn=transcribe_speech,\n",
    "    inputs=gr.Audio(sources=\"upload\", type=\"filepath\"),\n",
    "    outputs=gr.outputs.Textbox(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch the Gradio demo using our two blocks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with demo:\n",
    "    gr.TabbedInterface(\n",
    "        [mic_transcribe, file_transcribe],\n",
    "        [\"Transcribe Microphone\", \"Transcribe Audio File\"],\n",
    "    )\n",
    "\n",
    "demo.launch(debug=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}