notebooks/255-mms-massively-multilingual-speech/255-mms-massively-multilingual-speech.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c87087c91122e3f8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# MMS: Scaling Speech Technology to 1000+ languages with OpenVINO™\n",
    "\n",
    "The Massively Multilingual Speech (MMS) project expands speech technology from about 100 languages to over 1,000 by building a single multilingual speech recognition model supporting over 1,100 languages (more than 10 times as many as before), language identification models able to identify over 4,000 languages (40 times more than before), pretrained models supporting over 1,400 languages, and text-to-speech models for over 1,100 languages.\n",
    "The MMS model was proposed in [Scaling Speech Technology to 1,000+ Languages](https://arxiv.org/abs/2305.13516).  The models and code are originally released [here](https://github.com/facebookresearch/fairseq/tree/main/examples/mms).\n",
    "There are the different models open sourced in the MMS project: Automatic Speech Recognition (ASR), Language Identification (LID) and Speech Synthesis (TTS).  \n",
    "In this example we are considering ASR and LID. We will LID model to identify language, and then we will use ASR model corresponding to the language to recognize it. A simple diagram of this is below.\n",
    "![LID and ASR flow](https://github.com/openvinotoolkit/openvino_notebooks/assets/76171391/0e7fadd6-29a8-4fac-bd9c-41d66adcb045)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa80166a11177e7a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a id=\"0\"></a>\n",
    "### Table of contents:\n",
    "- [Prerequisites](#1)\n",
    "- [Prepare an example audio](#2)\n",
    "- [Language Identification (LID)](#3)\n",
    "  - [Download pretrained model and processor](#4)\n",
    "  - [Use the original model to run an inference](#5)\n",
    "  - [Convert to OpenVINO IR model and run an inference](#6)\n",
    "- [Automatic Speech Recognition (ASR)](#7)\n",
    "  - [Download pretrained model and processor](#8)\n",
    "  - [Use the original model to run an inference](#9)\n",
    "  - [Convert to OpenVINO IR model and run an inference](#10)\n",
    "- [Interactive demo with Gradio](#11)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c7a208b1fa497b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='1'></a>\n",
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc1a0304b8213aa",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%pip install -q --upgrade pip\n",
    "%pip install -q datasets \"transformers>=4.33.1\" accelerate \"openvino>=2023.1.0\" torch soundfile librosa \"numpy>=1.21.0,<=1.24\" gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbac6fae86122d9b",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "\n",
    "import openvino as ov"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d81ab16ec40431a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='2'></a>\n",
    "## Prepare an example audio\n",
    "Read an audio file and process the audio data. Make sure that the audio data is sampled to 16000 kHz.\n",
    "For this example we will use [a streamable version of the Multilingual LibriSpeech (MLS) dataset](https://huggingface.co/datasets/multilingual_librispeech). It supports contains example on 7 languages: `'german', 'dutch', 'french', 'spanish', 'italian', 'portuguese', 'polish'`.\n",
    "Choose one of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d46064f030034ef0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "\n",
    "SAMPLE_LANG = widgets.Dropdown(\n",
    "    options=['german', 'dutch', 'french', 'spanish', 'italian', 'portuguese', 'polish'],\n",
    "    value='german',\n",
    "    description='Dataset language:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "SAMPLE_LANG"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Specify `streaming=True` to not download the entire dataset."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "62f4f25bd4987849"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e3b30952e08ee76",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "mls = load_dataset(\"facebook/multilingual_librispeech\", SAMPLE_LANG.value, split=\"test\", streaming=True)\n",
    "mls = iter(mls)  # make it itarable\n",
    "\n",
    "example = next(mls)  # get one example"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68f9bb826d9a36dd",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Example has a dictionary structure. It contains an audio data and a text transcription."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d4ee3f9e30aacf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(example)  # look at structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "\n",
    "print(example['text'])\n",
    "ipd.Audio(example['audio']['array'], rate=16_000)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5f96bfecfd4bab51"
  },
  {
   "cell_type": "markdown",
   "id": "86963727a1d32e5a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='3'></a>\n",
    "## Language Identification (LID) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb607febc51e3782",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='4'></a>\n",
    "### Download pretrained model and processor\n",
    "Different LID models are available based on the number of languages they can recognize - 126, 256, 512, 1024, 2048, 4017. We will use 126."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1995f9336132be61",
   "metadata": {
    "collapsed": false,
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "from transformers import Wav2Vec2ForSequenceClassification, AutoFeatureExtractor\n",
    "\n",
    "model_id = \"facebook/mms-lid-126\"\n",
    "\n",
    "lid_processor = AutoFeatureExtractor.from_pretrained(model_id)\n",
    "lid_model = Wav2Vec2ForSequenceClassification.from_pretrained(model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "100d4f9dfff9a7d3",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='5'></a>\n",
    "### Use the original model to run an inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef184f78ef5f39c0",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "inputs = lid_processor(example['audio']['array'], sampling_rate=16_000, return_tensors=\"pt\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = lid_model(**inputs).logits\n",
    "\n",
    "lang_id = torch.argmax(outputs, dim=-1)[0].item()\n",
    "detected_lang = lid_model.config.id2label[lang_id]\n",
    "print(detected_lang)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bc6f53041bf77e4",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='6'></a>\n",
    "### Convert to OpenVINO IR model and run an inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8edf31c50d2c1c9a",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "MAX_SEQ_LENGTH = 30480\n",
    "\n",
    "input_values = torch.zeros([1, MAX_SEQ_LENGTH], dtype=torch.float)\n",
    "attention_mask = torch.zeros([1, MAX_SEQ_LENGTH], dtype=torch.int32)\n",
    "lid_model_xml_path = Path('models/ov_lid_model.xml')\n",
    "\n",
    "if not lid_model_xml_path.exists():\n",
    "    lid_model_xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    converted_model = ov.convert_model(lid_model, example_input={'input_values': input_values})\n",
    "    ov.save_model(converted_model, lid_model_xml_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a36953b5a21d5d6c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "And compile. Select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50ca1c1088501cf7",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "core = ov.Core()\n",
    "\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='AUTO',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "compiled_lid_model = core.compile_model(lid_model_xml_path, device_name=device.value)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "cf0b2d1c56693a83"
  },
  {
   "cell_type": "markdown",
   "id": "40193d2a396bb746",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Now it is possible to run an inference. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5d96a19f0504f3d",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def detect_lang(audio_data):\n",
    "    inputs = lid_processor(audio_data, sampling_rate=16_000, return_tensors=\"pt\")\n",
    "    \n",
    "    outputs = compiled_lid_model(inputs['input_values'])[0]\n",
    "    \n",
    "    lang_id = torch.argmax(torch.from_numpy(outputs), dim=-1)[0].item()\n",
    "    detected_lang = lid_model.config.id2label[lang_id]\n",
    "    \n",
    "    return detected_lang"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcaae46ecd2077b8",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "detect_lang(example['audio']['array'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "346a0954d96d40df",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Let's check another language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e89c90f-d6f0-4dc6-ba3d-34a4be710a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "SAMPLE_LANG = widgets.Dropdown(\n",
    "    options=['german', 'dutch', 'french', 'spanish', 'italian', 'portuguese', 'polish'],\n",
    "    value='french',\n",
    "    description='Dataset language:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "SAMPLE_LANG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e4b10f76be235ad",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mls = load_dataset(\"facebook/multilingual_librispeech\", SAMPLE_LANG.value, split=\"test\", streaming=True)\n",
    "mls = iter(mls)\n",
    "\n",
    "example = next(mls)\n",
    "print(example['text'])\n",
    "ipd.Audio(example['audio']['array'], rate=16_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f764403640f618",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "detect_language_id = detect_lang(example['audio']['array'])\n",
    "print(detect_language_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e010ed384d1e8ee7",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='7'></a>\n",
    "## Automatic Speech Recognition (ASR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe4536f63fe7e612",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='8'></a>\n",
    "### Download pretrained model and processor\n",
    "Download pretrained model and processor. By default, MMS loads adapter weights for English. If you want to load adapter weights of another language make sure to specify `target_lang=<your-chosen-target-lang>` as well as `ignore_mismatched_sizes=True`. The `ignore_mismatched_sizes=True` keyword has to be passed to allow the language model head to be resized according to the vocabulary of the specified language. Similarly, the processor should be loaded with the same target language. \n",
    "It is also possible to change the supported language later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b104f835667fb9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Wav2Vec2ForCTC, AutoProcessor\n",
    "model_id = \"facebook/mms-1b-all\"\n",
    "\n",
    "asr_processor = AutoProcessor.from_pretrained(model_id)\n",
    "asr_model = Wav2Vec2ForCTC.from_pretrained(model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5896f5fd08f62071",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "You can look at all supported languages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b62341511f98ceb",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "asr_processor.tokenizer.vocab.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "541c53d1c740d668",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Switch out the language adapters by calling the `load_adapter()` function for the model and `set_target_lang()` for the tokenizer. Pass the target language as an input - `\"detect_language_id\"` which was detected in the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15f9f4e31170b3fa",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "asr_processor.tokenizer.set_target_lang(detect_language_id)\n",
    "asr_model.load_adapter(detect_language_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de68b1eac717cc26",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='9'></a>\n",
    "### Use the original model to run an inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4463e26404e16195",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "inputs = asr_processor(example['audio']['array'], sampling_rate=16_000, return_tensors=\"pt\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = asr_model(**inputs).logits\n",
    "\n",
    "ids = torch.argmax(outputs, dim=-1)[0]\n",
    "transcription = asr_processor.decode(ids)\n",
    "print(transcription)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bda2f58170bfa2f4",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a name='10'></a>\n",
    "### Convert to OpenVINO IR model and run an inference\n",
    "Convert to OpenVINO IR model format with `ov.convert_model` function directly. Use `ov.save_model` function to serialize the result of conversion. For convenience of further use, we will create a function for these purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f47ccb726cdb505d",
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SEQ_LENGTH = 30480\n",
    "\n",
    "\n",
    "def get_asr_model(detected_language_id):\n",
    "    input_values = torch.zeros([1, MAX_SEQ_LENGTH], dtype=torch.float)\n",
    "    asr_model_xml_path = Path(f'models/ov_asr_{detected_language_id}_model.xml')\n",
    "    \n",
    "    if not asr_model_xml_path.exists():\n",
    "        asr_processor.tokenizer.set_target_lang(detected_language_id)\n",
    "        asr_model.load_adapter(detected_language_id)\n",
    "        \n",
    "        asr_model_xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "        converted_model = ov.convert_model(asr_model, example_input={'input_values': input_values})\n",
    "        ov.save_model(converted_model, asr_model_xml_path)\n",
    "    \n",
    "    compiled_asr_model = core.compile_model(asr_model_xml_path, device_name=device.value)\n",
    "\n",
    "    return compiled_asr_model\n",
    "\n",
    "\n",
    "compiled_asr_model = get_asr_model(detect_language_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4fb2cd466365800",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Run an inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83689739f10f2f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def recognize_audio(compiled_asr_model, src_audio):\n",
    "\n",
    "    inputs = asr_processor(src_audio, sampling_rate=16_000, return_tensors=\"pt\")\n",
    "    outputs = compiled_asr_model(inputs['input_values'])[0]\n",
    "    \n",
    "    ids = torch.argmax(torch.from_numpy(outputs), dim=-1)[0]\n",
    "    transcription = asr_processor.decode(ids)\n",
    "\n",
    "    return transcription\n",
    "\n",
    "\n",
    "transcription = recognize_audio(compiled_asr_model, example['audio']['array'])\n",
    "print(transcription)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "<a name='11'></a>\n",
    "## Interactive demo with Gradio\n",
    "In this demo you can try your own examples. Make sure that the audio data is sampled to 16000 kHz."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ba1a229fd290ec31"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "import librosa\n",
    "\n",
    "\n",
    "src_audio = gr.inputs.Audio(label=\"Source Audio\", type='filepath')\n",
    "outputs = [\n",
    "    gr.outputs.Textbox(label=\"Detected language ID\"),\n",
    "    gr.outputs.Textbox(label=\"Transcription\"),\n",
    "]\n",
    "title = 'MMS with Gradio'\n",
    "description = 'Gradio Demo for MMS and OpenVINO™. Upload a source audio, then click the \"Submit\" button to detect a language ID and a transcription. Make sure that the audio data is sampled to 16000 kHz. If this language has not been used before, it may take some time to prepare the ASR model.'\n",
    "\n",
    "\n",
    "def infer(src_audio_path):\n",
    "    src_audio, _ = librosa.load(src_audio_path)\n",
    "    detected_lang_id = detect_lang(src_audio)\n",
    "\n",
    "    yield detected_lang_id, None\n",
    "  \n",
    "    compiled_asr_model = get_asr_model(detected_lang_id)\n",
    "    transcription = recognize_audio(compiled_asr_model, src_audio)\n",
    "\n",
    "    yield detected_lang_id, transcription\n",
    "\n",
    "\n",
    "demo = gr.Interface(infer, src_audio, outputs, title=title, description=description)\n",
    "\n",
    "try:\n",
    "    demo.queue().launch(debug=True)\n",
    "except Exception:\n",
    "    demo.queue().launch(share=True, debug=True)\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "# demo.launch(server_name='your server name', server_port='server port in int')\n",
    "# Read more in the docs: https://gradio.app/docs/"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "49f4b7871f18ab46"
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}