431 lines
18 KiB
Python
431 lines
18 KiB
Python
"""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 datetime
|
|
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
|
|
WAIT_RETRY = 600 # seconds between tries while the built-in instance can't be reached
|
|
WAIT_MAX = 48 * 3600 # after this long since the button was pressed a waiting job gives up
|
|
WAIT_IDLE = 30 # seconds the runner sleeps while every queued job is waiting
|
|
|
|
_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
|
|
_not_before = {} # job id -> time.time() before which a waiting job isn't tried again
|
|
WAIT = 'wait' # _next_job: jobs are queued, but none is due yet
|
|
|
|
|
|
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 Unreachable(Exception):
|
|
"""The built-in instance can't be reached (e.g. the machine it runs on is off): the job waits
|
|
and tries again later instead of failing."""
|
|
|
|
|
|
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:
|
|
queued = list(TranscriptionJob.objects.filter(state=TranscriptionJob.QUEUED)
|
|
.order_by('id').values_list('id', flat=True))
|
|
if not queued:
|
|
_runner = None # decided under the lock, so a job queued right now starts a new runner
|
|
return None
|
|
now = time.time()
|
|
due = next((i for i in queued if _not_before.get(i, 0) <= now), None)
|
|
return due if due is not None else WAIT
|
|
|
|
|
|
def _loop():
|
|
global _runner
|
|
try:
|
|
while True:
|
|
close_old_connections()
|
|
job_id = _next_job()
|
|
if job_id is None:
|
|
return
|
|
if job_id == WAIT:
|
|
time.sleep(WAIT_IDLE)
|
|
continue
|
|
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 Unreachable as e:
|
|
_wait(job_id, str(e))
|
|
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 _wait(job_id, reason):
|
|
"""Back into the queue to try again in WAIT_RETRY seconds, or give up when it waited too long."""
|
|
job = TranscriptionJob.objects.filter(pk=job_id).only('created_at').first()
|
|
if job is None:
|
|
return
|
|
if (timezone.now() - job.created_at).total_seconds() > WAIT_MAX:
|
|
_not_before.pop(job_id, None)
|
|
_fail(job_id, 'Whisper-Instanz zwei Tage lang nicht erreichbar')
|
|
return
|
|
_not_before[job_id] = time.time() + WAIT_RETRY
|
|
again = timezone.localtime(timezone.now() + datetime.timedelta(seconds=WAIT_RETRY)).strftime('%H:%M')
|
|
log.info('transcription job %s waits for the instance (%s)', job_id, reason)
|
|
TranscriptionJob.objects.filter(pk=job_id).update(
|
|
state=TranscriptionJob.QUEUED, chunks_done=0, chunks_total=0, started_at=None,
|
|
error=f'Whisper-Instanz nicht erreichbar, neuer Versuch um {again}')
|
|
|
|
|
|
def _probe(url):
|
|
"""Unreachable unless something answers at the built-in instance's address."""
|
|
try:
|
|
requests.get(url.rsplit('/audio/transcriptions', 1)[0], timeout=(5, 10)).close()
|
|
except (requests.ConnectionError, requests.Timeout) as e:
|
|
raise Unreachable(e.__class__.__name__)
|
|
|
|
|
|
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)
|
|
if builtin:
|
|
_probe(url) # before downloading: is the instance there at all?
|
|
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.
|
|
try:
|
|
r = requests.post(url, headers=headers, files=files, data=options, timeout=UPLOAD_TIMEOUT)
|
|
except requests.ConnectionError as e:
|
|
raise Unreachable(e.__class__.__name__)
|
|
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 Unreachable:
|
|
if attempt >= RETRIES:
|
|
raise
|
|
time.sleep(RETRY_PAUSES[min(attempt, len(RETRY_PAUSES) - 1)])
|
|
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}}
|