core-transcription/speaker_labelled_subtitles.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Creating Subtitles with Speaker Labels"
      ],
      "metadata": {
        "id": "9ie0NWRRzFwp"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This Colab will demonstrate how to use AssemblyAI's [Speaker Diarization](https://www.assemblyai.com/docs/speech-to-text/speaker-diarization) model together to format subtitles according to their respective speaker.\n",
        "\n",
        "# Getting Started\n",
        "Before we begin, make sure you have an AssemblyAI account and an API key. You can [sign up](https://www.assemblyai.com/dashboard/signup) for an AssemblyAI account and get your API key from your [dashboard](https://www.assemblyai.com/app/account).\n",
        "\n"
      ],
      "metadata": {
        "id": "Dwd9Zxmo4olH"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pc6nDAMqzBDE"
      },
      "outputs": [],
      "source": [
        "!pip install assemblyai"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, we will configure our API key as well as our file to be transcribed. Then, we decide on a number of words we want to have per subtitle.\n",
        "\n",
        "Lastly, we transcribe our file."
      ],
      "metadata": {
        "id": "zFGsz9YlGIlQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import assemblyai as aai\n",
        "\n",
        "# SETTINGS\n",
        "aai.settings.api_key = \"YOUR-API-KEY\"\n",
        "filename = \"YOUR-FILE-NAME\"\n",
        "transcriber = aai.Transcriber(config=aai.TranscriptionConfig(speaker_labels=True))\n",
        "transcript = transcriber.transcribe(filename)\n",
        "\n",
        "# Maximum number of words per subtitle\n",
        "max_words_per_subtitle = 6"
      ],
      "metadata": {
        "id": "FnJfmd174yyo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## How the code works\n",
        "`speaker_colors` is a dictionary that maps speaker identifiers (like \"A\", \"B\", \"C\", etc.) to specific colors. Each speaker in the transcription will be associated with a unique color in the subtitles.\n",
        "\n",
        "When Speaker Diarization is enabled, sentences in our API response have a speaker code under the `speaker` key. We use the speaker code to determine the color of the subtitle text."
      ],
      "metadata": {
        "id": "YEQu3RJgImp3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Color assignments for speakers\n",
        "speaker_colors = {\n",
        "    \"A\": \"red\",\n",
        "    \"B\": \"orange\",\n",
        "    \"C\": \"yellow\",\n",
        "    \"D\": \"yellowgreen\",\n",
        "    \"E\": \"green\",\n",
        "    \"F\": \"lightskyblue\",\n",
        "    \"G\": \"purple\",\n",
        "    \"H\": \"mediumpurple\",\n",
        "    \"I\": \"pink\",\n",
        "    \"J\": \"brown\",\n",
        "}\n",
        "\n",
        "# Process transcription segments\n",
        "def process_segments(segments):\n",
        "    srt_content = \"\"\n",
        "    subtitle_index = 1\n",
        "    for segment in segments:\n",
        "        speaker = segment.speaker\n",
        "        color = speaker_colors.get(speaker, \"black\")  # Default color is black\n",
        "\n",
        "        # Split text into words and group into chunks\n",
        "        words = segment.words\n",
        "        for i in range(0, len(words), max_words_per_subtitle):\n",
        "            chunk = words[i:i + max_words_per_subtitle]\n",
        "            start_time = chunk[0].start  # -1 indicates continuation\n",
        "            end_time = chunk[-1].end\n",
        "            srt_content += create_subtitle(subtitle_index, start_time, end_time, chunk, color)\n",
        "            subtitle_index += 1\n",
        "\n",
        "    return srt_content\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Create a single subtitle\n",
        "def create_subtitle(index, start_time, end_time, words, color):\n",
        "    \"\"\"\n",
        "    If you're previewing this code on GitHub, take note that this function generates HTML for color styling,\n",
        "    and this HTML is not rendered in the GitHub preview. Run the notebook locally to see the actual code.\n",
        "    Visit this link to learn more about formatting SRT files: https://docs.fileformat.com/video/srt/#formatting-of-srt-files\n",
        "    \"\"\"\n",
        "    text = \"\"\n",
        "    for word in words:\n",
        "        text += word.text + ' '\n",
        "    start_srt = format_time(start_time)\n",
        "    end_srt = format_time(end_time)\n",
        "    return f\"{index}\\n{start_srt} --> {end_srt}\\n<font color=\\\"{color}\\\">{text}</font>\\n\\n\"\n",
        "\n",
        "# Format time in SRT style\n",
        "def format_time(milliseconds):\n",
        "    hours, remainder = divmod(milliseconds, 3600000)\n",
        "    minutes, remainder = divmod(remainder, 60000)\n",
        "    seconds, milliseconds = divmod(remainder, 1000)\n",
        "    return f\"{int(hours):02}:{int(minutes):02}:{int(seconds):02},{int(milliseconds):03}\""
      ],
      "metadata": {
        "id": "gkOsP0KAFmEY"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Our last step is to generate and save our subtitle file!"
      ],
      "metadata": {
        "id": "4guym9y3KZZ3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate SRT content\n",
        "sentences = transcript.get_sentences()\n",
        "srt_content = process_segments(sentences)\n",
        "\n",
        "# Save to SRT file\n",
        "with open(filename + '.srt', 'w') as file:\n",
        "    file.write(srt_content)\n",
        "\n",
        "print(f\"SRT file generated: {filename}.srt\")"
      ],
      "metadata": {
        "id": "0rXnMeJWFsF1"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}