diora-web/podcasts/transcribe.py

"""Transcribe an episode with a Whisper instance (docs/transcription.md, "Transcribing").

The server downloads the episode (as silence.py does, with the same size limit), decodes it
once to 16 kHz mono, cuts it into pieces of about eight minutes where there is a pause (so no
word is cut), sends the pieces in parallel to the instance's OpenAI-compatible
`/audio/transcriptions` and joins the answers into one transcript, which is kept on the job
(TranscriptionJob.data).

A job runs in a daemon thread, one at a time per server; the others wait as `queued` and are
taken oldest first (by id). A job found `downloading` or `transcribing` when the process has
just started (a restart killed it) is queued again. Cancelling deletes the job's row: the
worker notices when it next writes progress and stops.
"""
import logging
import os
import subprocess
import tempfile
import threading
import time
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait

import requests
from django.conf import settings
from django.db import close_old_connections
from django.utils import timezone

from . import chapters, silence, transcripts
from .models import TranscriptionJob, TranscriptionSettings
from .transcription import DEFAULT_MODEL, for_user, token_of

log = logging.getLogger(__name__)

CHUNK_SECONDS = 8 * 60      # about this long a piece
CUT_WINDOW = 30             # a cut moves to the longest pause within this many seconds
MIN_TAIL = 60               # a last piece shorter than this is joined to the one before
RATE = 16000                # Hz, mono, 16-bit
WINDOW = silence.WINDOW
RETRIES = 3
RETRY_PAUSES = (3, 10, 30)  # seconds before each retry
UPLOAD_TIMEOUT = (10, 3600)  # a CPU instance may need many minutes for a piece
DECODE_TIMEOUT = 3600
OPUS_BITRATE = '32k'
MAX_PROMPT = 500            # characters
ALIVE_CHECK = 5             # seconds between looks whether the job was cancelled

_lock = threading.Lock()
_runner = None              # the thread working through the queue, or None
_resumed = False            # whether jobs left over from before a restart were queued again


class Cancelled(Exception):
    """The job's row is gone: the user cancelled it."""


class Failed(Exception):
    """The job can't succeed; the message is shown to the user."""


class Permanent(Exception):
    """A request the instance refused for good (a client error): not worth a retry."""


# ---------------------------------------------------------------------------
# Queue
# ---------------------------------------------------------------------------

def kick():
    """Make sure queued jobs are being worked on. Called by the endpoints, so after a restart
    the first request brings back what was running."""
    global _runner, _resumed
    with _lock:
        if not _resumed:
            _resumed = True
            TranscriptionJob.objects.filter(
                state__in=[TranscriptionJob.DOWNLOADING, TranscriptionJob.TRANSCRIBING],
            ).update(state=TranscriptionJob.QUEUED, chunks_done=0, chunks_total=0)
        if _runner is not None:
            return
        if not TranscriptionJob.objects.filter(state=TranscriptionJob.QUEUED).exists():
            return
        _runner = threading.Thread(target=_loop, name='transcribe', daemon=True)
        _runner.start()


def _next_job():
    global _runner
    with _lock:
        job_id = (TranscriptionJob.objects.filter(state=TranscriptionJob.QUEUED)
                  .order_by('id').values_list('id', flat=True).first())
        if job_id is None:
            _runner = None   # decided under the lock, so a job queued right now starts a new runner
        return job_id


def _loop():
    global _runner
    try:
        while True:
            close_old_connections()
            job_id = _next_job()
            if job_id is None:
                return
            try:
                run(job_id)
            except Exception:  # run() reports its own failures; this keeps the queue going
                log.exception('transcription job %s crashed', job_id)
    except BaseException:
        with _lock:
            _runner = None
        raise
    finally:
        close_old_connections()


def run(job_id):
    """Work one job through to `done` or `error`, in this thread."""
    started = TranscriptionJob.objects.filter(pk=job_id, state=TranscriptionJob.QUEUED).update(
        state=TranscriptionJob.DOWNLOADING, started_at=timezone.now(), finished_at=None,
        error='', chunks_done=0, chunks_total=0)
    if not started:
        return   # cancelled while it waited
    try:
        data = _transcribe(job_id)
        _write(job_id, state=TranscriptionJob.DONE, data=data, error='', finished_at=timezone.now())
    except Cancelled:
        log.info('transcription job %s cancelled', job_id)
    except Failed as e:
        _fail(job_id, str(e))
    except Exception as e:
        log.exception('transcription job %s failed', job_id)
        _fail(job_id, f'{type(e).__name__}: {e}')


def _fail(job_id, message):
    TranscriptionJob.objects.filter(pk=job_id).update(
        state=TranscriptionJob.ERROR, error=(message or 'failed')[:500], finished_at=timezone.now())


def _write(job_id, **fields):
    """Update the job; Cancelled when its row is gone."""
    if not TranscriptionJob.objects.filter(pk=job_id).update(**fields):
        raise Cancelled()


def _alive(job_id):
    if not TranscriptionJob.objects.filter(pk=job_id).exists():
        raise Cancelled()


# ---------------------------------------------------------------------------
# The work
# ---------------------------------------------------------------------------

def _target(conf):
    """(url, headers, builtin) of the instance the settings point at."""
    builtin = conf.engine == TranscriptionSettings.BUILTIN
    base = settings.WHISPER_URL if builtin else conf.url
    if not base:
        raise Failed('Keine eingebaute Instanz eingerichtet' if builtin else 'Keine Adresse eingetragen')
    headers = {'User-Agent': 'diora transcription'}
    if not builtin and conf.token_encrypted:
        headers['Authorization'] = 'Bearer ' + token_of(conf)
    return base.rstrip('/') + '/audio/transcriptions', headers, builtin


def _prompt(episode):
    parts = [episode.feed.title, episode.title, episode.feed.author]
    return '. '.join(p.strip() for p in parts if p and p.strip())[:MAX_PROMPT]


def _transcribe(job_id):
    job = TranscriptionJob.objects.select_related('episode__feed').get(pk=job_id)
    episode = job.episode
    conf = for_user(episode.feed.user)
    url, headers, builtin = _target(conf)
    options = {'model': conf.model or DEFAULT_MODEL, 'response_format': 'verbose_json'}
    if conf.language:
        options['language'] = conf.language
    prompt = _prompt(episode)
    if prompt:
        options['prompt'] = prompt

    with tempfile.TemporaryDirectory(prefix='diora-transcribe-') as tmp:
        source = os.path.join(tmp, 'episode.audio')
        _download(episode.audio_url, source, job_id)
        pcm = os.path.join(tmp, 'episode.pcm')
        _decode(source, pcm)
        os.remove(source)
        total = os.path.getsize(pcm) // 2 / RATE
        if total <= 0:
            raise Failed('Kein Ton in der Datei')
        pieces = plan(total, pauses_of(pcm))
        _write(job_id, state=TranscriptionJob.TRANSCRIBING, chunks_total=len(pieces), chunks_done=0)

        results = {}
        stop = threading.Event()   # set when the job ends early: no piece starts after that
        suffix, mime = ('.flac', 'audio/flac') if builtin else ('.ogg', 'audio/ogg')

        def work(index):
            # Only the HTTP side runs here; the database is touched by this thread alone.
            if stop.is_set():
                return index, None
            start, end = pieces[index]
            path = _encode(pcm, start, end, tmp, index, opus=not builtin)
            try:
                with open(path, 'rb') as f:
                    body = f.read()
            finally:
                os.remove(path)
            return index, _send(url, headers, builtin, f'piece-{index}{suffix}', mime, body, options, stop)

        pool = ThreadPoolExecutor(max_workers=max(1, min(conf.parallel, 8)))
        try:
            pending = {pool.submit(work, i) for i in range(len(pieces))}
            while pending:
                finished, pending = wait(pending, timeout=ALIVE_CHECK, return_when=FIRST_COMPLETED)
                _alive(job_id)   # also while a long piece is out
                for f in finished:
                    index, answer = f.result()   # the first failure ends the job
                    results[index] = answer
                if finished:
                    _write(job_id, chunks_done=len(results))
        finally:
            stop.set()
            pool.shutdown(wait=True, cancel_futures=True)
    return build(pieces, [results[i] for i in range(len(pieces))], options['model'], conf.language)


def _download(url, path, job_id):
    size = 0
    with open(path, 'wb') as out, requests.get(
            url, stream=True, timeout=(10, 60), headers={'User-Agent': 'diora transcription'}) as r:
        if r.status_code >= 400:
            raise Failed(f'Download: HTTP {r.status_code}')
        if int(r.headers.get('Content-Length') or 0) > silence.MAX_BYTES:
            raise Failed('too large')
        for n, chunk in enumerate(r.iter_content(1 << 20)):
            size += len(chunk)
            if size > silence.MAX_BYTES:
                raise Failed('too large')
            out.write(chunk)
            if n % 8 == 0:
                _alive(job_id)
    if size == 0:
        raise Failed('Download: leere Datei')


def _run_ffmpeg(args, input=None, timeout=DECODE_TIMEOUT):
    import imageio_ffmpeg
    proc = subprocess.run(
        [imageio_ffmpeg.get_ffmpeg_exe(), '-nostdin', '-hide_banner', '-loglevel', 'error', *args],
        input=input, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout)
    if proc.returncode != 0:
        err = proc.stderr.decode('utf-8', 'replace').strip().splitlines()
        raise Failed('ffmpeg: ' + (err[-1] if err else f'exit {proc.returncode}')[:200])


def _decode(source, pcm):
    """The audio at `source` as raw 16 kHz mono 16-bit samples in the file `pcm`."""
    _run_ffmpeg(['-vn', '-sn', '-dn', '-i', source, '-ac', '1', '-ar', str(RATE), '-f', 's16le', '-y', pcm])


def pauses_of(pcm):
    """The pauses of the decoded audio as [(middle, length)] in seconds, found with
    silence.py's rules (what it cuts out of a pause is centred on the pause)."""
    size = int(RATE * WINDOW) * 2
    levels = []
    with open(pcm, 'rb') as f:
        while True:
            data = f.read(size * 100000)
            if not data:
                break
            for i in range(0, len(data) - size + 1, size):
                levels.append(silence._level_db(data[i:i + size]))
    return [((a + b) / 2, b - a) for a, b in silence.cuts(levels)]


def plan(total, pauses):
    """[(start, end)] in seconds covering 0..total: pieces of about CHUNK_SECONDS, each cut at
    the longest pause within CUT_WINDOW of where it would fall (right there when there is none)."""
    cuts = []
    last = 0.0
    k = 1
    while k * CHUNK_SECONDS < total - MIN_TAIL:
        target = k * CHUNK_SECONDS
        k += 1
        near = [p for p in pauses if abs(p[0] - target) <= CUT_WINDOW and p[0] > last + 1]
        at = max(near, key=lambda p: p[1])[0] if near else target
        if at <= last or at >= total:
            continue
        cuts.append(at)
        last = at
    edges = [0.0, *cuts, total]
    return [(edges[i], edges[i + 1]) for i in range(len(edges) - 1)]


def _encode(pcm, start, end, tmp, index, opus):
    """The piece [start, end) of the decoded audio as FLAC or Opus (ogg), in a file."""
    with open(pcm, 'rb') as f:
        f.seek(int(round(start * RATE)) * 2)
        data = f.read(int(round((end - start) * RATE)) * 2)
    out = os.path.join(tmp, f'piece-{index}.' + ('ogg' if opus else 'flac'))
    codec = ['-c:a', 'libopus', '-b:a', OPUS_BITRATE, '-f', 'ogg'] if opus else ['-c:a', 'flac', '-f', 'flac']
    _run_ffmpeg(['-f', 's16le', '-ar', str(RATE), '-ac', '1', '-i', 'pipe:0', *codec, '-y', out], input=data)
    return out


def _post(url, headers, builtin, name, mime, body, options):
    """One request; the parsed answer. Permanent for a refusal that won't change, RuntimeError
    for trouble that may pass."""
    files = {'file': (name, body, mime)}
    try:
        if builtin:
            # The operator's own instance (DIORA_WHISPER_URL) may sit on a private network.
            r = requests.post(url, headers=headers, files=files, data=options, timeout=UPLOAD_TIMEOUT)
            if r.status_code >= 400:
                if r.status_code < 500 and r.status_code not in (408, 429):
                    raise Permanent(f'HTTP {r.status_code}')
                raise RuntimeError(f'HTTP {r.status_code}')
        else:
            try:
                r = chapters._open(url, headers, method='POST', timeout=UPLOAD_TIMEOUT, files=files, data=options)
            except chapters.ChaptersInvalid as e:
                raise Permanent(str(e))
            except chapters.ChaptersUnavailable as e:
                raise RuntimeError(str(e))
        try:
            return r.json()
        finally:
            r.close()
    except requests.RequestException as e:
        raise RuntimeError(e.__class__.__name__)
    except ValueError:
        raise RuntimeError('Antwort ist kein JSON')


def _send(url, headers, builtin, name, mime, body, options, stop):
    """A piece to the instance, retried up to RETRIES times with a growing pause."""
    for attempt in range(RETRIES + 1):
        if stop.is_set():
            raise Cancelled()
        try:
            return _post(url, headers, builtin, name, mime, body, options)
        except Permanent as e:
            raise Failed(f'Whisper: {e}')
        except RuntimeError as e:
            if attempt >= RETRIES:
                raise Failed(f'Whisper: {e}')
            log.info('whisper piece %s failed (%s), retrying', name, e)
            time.sleep(RETRY_PAUSES[min(attempt, len(RETRY_PAUSES) - 1)])


def build(pieces, answers, model, language):
    """{segments, source} from the pieces' places and the instance's verbose_json answers."""
    fragments = []
    detected = ''
    for (start, end), answer in zip(pieces, answers):
        if not isinstance(answer, dict):
            continue
        detected = detected or (answer.get('language') or '')
        segments = answer.get('segments')
        if not isinstance(segments, list):   # a plain {"text": …} answer: one segment for the piece
            text = (answer.get('text') or '').strip()
            segments = [{'start': 0, 'end': end - start, 'text': text}] if text else []
        for s in segments:
            if not isinstance(s, dict):
                continue
            text = ' '.join(str(s.get('text') or '').split())
            try:
                a, b = float(s.get('start', 0)), float(s.get('end', 0))
            except (TypeError, ValueError):
                continue
            if text:
                fragments.append((start + a, start + b, '', text))
    merged = transcripts.merge(fragments)
    if not merged:
        raise Failed('Kein Text erkannt')
    return {'segments': merged, 'source': {'type': 'whisper', 'model': model, 'language': language or detected}}