core-transcription/subtitle_creation_by_word_count.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": [
        "# Create custom length subtitles with AssemblyAI\n",
        "\n"
      ],
      "metadata": {
        "id": "6jpYYU6Eh14M"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "While our SRT/VTT endpoints do allow you to customize the maximum number of characters per caption using the chars_per_caption URL parameter in your API requests, there are some use-cases that require a custom number of words in each subtitle.\n",
        "\n",
        "In this Colab, we will demonstrate how to construct these subtitles yourself in Python!"
      ],
      "metadata": {
        "id": "lOkGwZOeit5E"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step-by-Step Instructions"
      ],
      "metadata": {
        "id": "ejk9-USqjVpF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install -U assemblyai\n"
      ],
      "metadata": {
        "id": "8K4rTtDAjefo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import the assemblyai package and set the API key.\n",
        "\n"
      ],
      "metadata": {
        "id": "MwwCHI1AjYdJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import assemblyai as aai\n",
        "\n",
        "aai.settings.api_key = \"YOUR-API-KEY\""
      ],
      "metadata": {
        "id": "AyE7BUBnjfkt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create a Transcriber object.\n"
      ],
      "metadata": {
        "id": "vyDXhkqkjhXh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "transcriber = aai.Transcriber()"
      ],
      "metadata": {
        "id": "xrx_D4cxjjfh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Use the Transcriber object's transcribe method and pass in the audio file's path as a parameter. The transcribe method saves the results of the transcription to the Transcriber object's transcript attribute."
      ],
      "metadata": {
        "id": "0Tnxu229jpDE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "transcript = transcriber.transcribe(\"./my-audio.mp3\")"
      ],
      "metadata": {
        "id": "yghza_7tjrgQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Alternatively, you can pass in the URL of the publicly accessible audio file on the internet.\n",
        "\n"
      ],
      "metadata": {
        "id": "dkOKPKBSjtA8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "transcript = transcriber.transcribe(\"https://storage.googleapis.com/aai-docs-samples/espn.m4a\")"
      ],
      "metadata": {
        "id": "dhZvEAl2joyF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a function that converts seconds to timecodes"
      ],
      "metadata": {
        "id": "4a-Dbw8fj5x9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def second_to_timecode(x: float) -> str:\n",
        "    hour, x = divmod(x, 3600)\n",
        "    minute, x = divmod(x, 60)\n",
        "    second, x = divmod(x, 1)\n",
        "    millisecond = int(x * 1000.)\n",
        "\n",
        "    return '%.2d:%.2d:%.2d,%.3d' % (hour, minute, second, millisecond)"
      ],
      "metadata": {
        "id": "BdXWruwNkAZH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a function that iterates through the transcripts object to construct a list according to the number of words per subtitle"
      ],
      "metadata": {
        "id": "N6sZzPD7kE26"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BHtOhJqkh0KI"
      },
      "outputs": [],
      "source": [
        "def generate_subtitles_by_word_count(transcript, words_per_line):\n",
        "  output = []\n",
        "  subtitle_index = 1  # Start subtitle index at 1\n",
        "  word_count = 0\n",
        "  current_words = []\n",
        "\n",
        "  for sentence in transcript.get_sentences():\n",
        "    for word in sentence.words:\n",
        "      current_words.append(word)\n",
        "      word_count += 1\n",
        "      if word_count >= words_per_line or word == sentence.words[-1]:\n",
        "        start_time = second_to_timecode(current_words[0].start / 1000)\n",
        "        end_time = second_to_timecode(current_words[-1].end / 1000)\n",
        "        subtitle_text = \" \".join([word.text for word in current_words])\n",
        "        output.append(str(subtitle_index))\n",
        "        output.append(\"%s --> %s\" % (start_time, end_time))\n",
        "        output.append(subtitle_text)\n",
        "        output.append(\"\")\n",
        "        current_words = []  # Reset for the next subtitle\n",
        "        word_count = 0  # Reset word count\n",
        "        subtitle_index += 1\n",
        "\n",
        "  return output"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate your subtitle file"
      ],
      "metadata": {
        "id": "n-3pqh9ClGUn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "subs = generate_subtitles_by_word_count(transcript, 6)\n",
        "with open(f\"{transcript.id}.srt\", 'w') as o:\n",
        "    final = '\\n'.join(subs)\n",
        "    o.write(final)\n",
        "\n",
        "print(\"SRT file generated.\")"
      ],
      "metadata": {
        "id": "T8pOkDjQlOUi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Conclusion\n",
        "This is a useful script for those looking to customise the number of words in their captions instead of the number of characters!"
      ],
      "metadata": {
        "id": "Bo2TwNyOnd7F"
      }
    }
  ]
}