Transcribe audio to SRT/VTT captions

import re

from fishaudio import FishAudio
from fishaudio.core import RequestOptions

client = FishAudio()

MAX_CHARS = 42  # longest cue text, in characters
MAX_MS = 6000  # longest cue, in milliseconds
MAX_GAP_MS = 600  # a longer pause between words starts a new cue
MIN_MS = 500  # display time for a cue whose words have no duration

# Chinese and Japanese are written without spaces between words.
NO_SPACE = re.compile(r"[\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff]")


def join_words(left: str, right: str) -> str:
    if NO_SPACE.match(left[-1]) and NO_SPACE.match(right[0]):
        return left + right
    return f"{left} {right}"


def group_cues(segments) -> list[dict]:
    """Group word-level segments into cues, with start and end in milliseconds."""
    cues = []
    for segment in segments:
        word = segment.text.strip()
        if not word:
            continue
        start, end = round(segment.start * 1000), round(segment.end * 1000)
        if cues:
            cue = cues[-1]
            text = join_words(cue["text"], word)
            if (
                start - cue["end"] <= MAX_GAP_MS
                and len(text) <= MAX_CHARS
                and end - cue["start"] <= MAX_MS
            ):
                cue["text"], cue["end"] = text, max(cue["end"], end)
                continue
        cues.append({"start": start, "end": end, "text": word})

    # A segment can have start == end. Give a cue with no duration a short one,
    # ending no later than the next cue starts, so every cue ends after it starts.
    for i, cue in enumerate(cues):
        if cue["end"] <= cue["start"]:
            end = cue["start"] + MIN_MS
            if i + 1 < len(cues):
                end = min(end, cues[i + 1]["start"])
            cue["end"] = max(end, cue["start"] + 1)
    return cues


def format_timestamp(ms: int, separator: str) -> str:
    """HH:MM:SS,mmm for SRT (separator ","), HH:MM:SS.mmm for WebVTT (".")."""
    hours, ms = divmod(ms, 3_600_000)
    minutes, ms = divmod(ms, 60_000)
    seconds, ms = divmod(ms, 1000)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}{separator}{ms:03d}"


with open("speech.wav", "rb") as f:
    result = client.asr.transcribe(
        audio=f.read(),
        include_timestamps=True,
        request_options=RequestOptions(
            timeout=900,  # long recordings can take several minutes
            additional_headers={"model": "transcribe-1-pro"},
        ),
    )

cues = group_cues(result.segments)

# SRT: 1-based index, comma decimal separator, blank line between cues.
with open("captions.srt", "w", encoding="utf-8") as srt:
    for i, cue in enumerate(cues, start=1):
        start = format_timestamp(cue["start"], ",")
        end = format_timestamp(cue["end"], ",")
        srt.write(f"{i}\n{start} --> {end}\n{cue['text']}\n\n")

# WebVTT: same cues, "WEBVTT" header, dot decimal separator.
with open("captions.vtt", "w", encoding="utf-8") as vtt:
    vtt.write("WEBVTT\n\n")
    for cue in cues:
        start = format_timestamp(cue["start"], ".")
        end = format_timestamp(cue["end"], ".")
        vtt.write(f"{start} --> {end}\n{cue['text']}\n\n")

print(f"Wrote {len(cues)} cues to captions.srt and captions.vtt")

import { FishAudioClient } from "fish-audio";
import { readFile, writeFile } from "fs/promises";

const client = new FishAudioClient({ apiKey: process.env.FISH_API_KEY });

const MAX_CHARS = 42; // longest cue text, in characters
const MAX_MS = 6000; // longest cue, in milliseconds
const MAX_GAP_MS = 600; // a longer pause between words starts a new cue
const MIN_MS = 500; // display time for a cue whose words have no duration

// Chinese and Japanese are written without spaces between words.
const NO_SPACE = /[\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff]/;

function joinWords(left, right) {
  if (NO_SPACE.test(left.at(-1)) && NO_SPACE.test(right[0])) {
    return left + right;
  }
  return `${left} ${right}`;
}

// Group word-level segments into cues, with start and end in milliseconds.
function groupCues(segments) {
  const cues = [];
  for (const segment of segments) {
    const word = segment.text.trim();
    if (!word) continue;
    const start = Math.round(segment.start * 1000);
    const end = Math.round(segment.end * 1000);
    const cue = cues.at(-1);
    if (cue) {
      const text = joinWords(cue.text, word);
      if (
        start - cue.end <= MAX_GAP_MS &&
        text.length <= MAX_CHARS &&
        end - cue.start <= MAX_MS
      ) {
        cue.text = text;
        cue.end = Math.max(cue.end, end);
        continue;
      }
    }
    cues.push({ start, end, text: word });
  }

  // A segment can have start == end. Give a cue with no duration a short one,
  // ending no later than the next cue starts, so every cue ends after it starts.
  cues.forEach((cue, i) => {
    if (cue.end <= cue.start) {
      let end = cue.start + MIN_MS;
      if (i + 1 < cues.length) end = Math.min(end, cues[i + 1].start);
      cue.end = Math.max(end, cue.start + 1);
    }
  });
  return cues;
}

// HH:MM:SS,mmm for SRT (separator ","), HH:MM:SS.mmm for WebVTT (".").
function formatTimestamp(ms, separator) {
  const pad = (n, width) => String(n).padStart(width, "0");
  const hours = Math.floor(ms / 3_600_000);
  const minutes = Math.floor((ms % 3_600_000) / 60_000);
  const seconds = Math.floor((ms % 60_000) / 1000);
  const millis = ms % 1000;
  return `${pad(hours, 2)}:${pad(minutes, 2)}:${pad(seconds, 2)}${separator}${pad(millis, 3)}`;
}

const result = await client.speechToText.convert(
  {
    audio: new File([await readFile("speech.wav")], "speech.wav"),
    ignore_timestamps: false,
  },
  {
    headers: { model: "transcribe-1-pro" },
    timeoutInSeconds: 900, // long recordings can take several minutes
  }
);

const cues = groupCues(result.segments);
const timing = (cue, separator) =>
  `${formatTimestamp(cue.start, separator)} --> ${formatTimestamp(cue.end, separator)}`;

// SRT: 1-based index, comma decimal separator, blank line between cues.
const srt = cues.map(
  (cue, i) => `${i + 1}\n${timing(cue, ",")}\n${cue.text}\n\n`
);
await writeFile("captions.srt", srt.join(""), "utf-8");

// WebVTT: same cues, "WEBVTT" header, dot decimal separator.
const vtt = cues.map(cue => `${timing(cue, ".")}\n${cue.text}\n\n`);
await writeFile("captions.vtt", "WEBVTT\n\n" + vtt.join(""), "utf-8");

console.log(`Wrote ${cues.length} cues to captions.srt and captions.vtt`);