danieledll/moss-transcribe-diarize-atc

MOSS-Transcribe-Diarize — ATC fine-tune

Fine-tune of OpenMOSS-Team/MOSS-Transcribe-Diarize (0.9B, joint transcription + speaker diarization) specialized for US air-traffic-control radio audio (VHF AM feeds, ~30 concurrent voices per 30-minute clip).

Training data

Publicly available ATC audio and recordings: stitched multi-speaker training samples with controlled voice recurrence, plus ~230 h of self-labeled, human-reviewed recordings across 35+ US airports (ground / tower / approach / departure), with diarization labels anchored to ADS-B-derived callsign identities.

Training compute

Iteratively fine-tuned on a single GPU: early iterations on 1× H100 80 GB, later ones on 1× H200 141 GB — the larger memory raised the usable sample-length cap, letting full-length (~30-minute, up to ~40k-token) clips join the training mix instead of being filtered out. Total fine-tuning compute across all iterations: roughly $250 of on-demand GPU time.

Evaluation (vs the base model)

benchmark	base	fine-tuned
32-min real tower clip — speaker purity	0.35 (8.5/44 voices found)	0.930 (41/41 identities)
10-clip unseen-airport eval — pair-F1 / greedy-acc	0.457 / 0.593 (5-clip subset)	0.644 / 0.719
US tower audio — WER vs adjudicated gold	0.162	0.115
Held-out stitched gold test — WER	0.551	0.427

1. Input audio: the denoise chain (REQUIRED)

The model was trained exclusively on audio processed by this exact ffmpeg chain, and it under-performs on raw un-denoised audio — run every input through it first:

ffmpeg -i input.mp3 \
  -af "highpass=f=200,lowpass=f=3600,afftdn=nr=12:nf=-25,dynaudnorm=f=150:g=15" \
  -ar 16000 -ac 1 -sample_fmt s16 input.wav

(High-pass 200 Hz and low-pass 3600 Hz bracket the AM voice band; afftdn removes broadband hiss; dynaudnorm levels the wildly varying per-transmission loudness of scanner feeds.)

2. Loading (note the processor step)

Load the weights from this repository, but the processor from the base model at the pinned revision: AutoProcessor pointed at a fine-tune checkpoint silently falls back to a bare text tokenizer (the saved config carries no processor auto_map), which breaks audio featurization without an error.

import soundfile as sf
import torch
from transformers import AutoModelForCausalLM, AutoProcessor

REPO = "danieledll/moss-transcribe-diarize-atc"
BASE = "OpenMOSS-Team/MOSS-Transcribe-Diarize"
BASE_REV = "e8681d68e7042738ffca8ac8212bc8fcb1131ab8"

PROMPT = (
    "Transcribe the audio. For each segment, start with the timestamp and speaker ID "
    "([S01], [S02], [S03], ...), then the spoken text, and end with the segment timestamp."
)  # the exact prompt the checkpoint was trained with — do not change it

model = (
    AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True, dtype="auto")
    .to(torch.bfloat16).to("cuda").eval()
)
processor = AutoProcessor.from_pretrained(BASE, revision=BASE_REV, trust_remote_code=True)

3. Transcription

audio, sr = sf.read("input.wav", dtype="float32", always_2d=True)
assert sr == 16000
audio = audio.mean(axis=1)  # mono

messages = [{"role": "user", "content": [
    {"type": "audio", "audio": ""}, {"type": "text", "text": PROMPT},
]}]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    inputs = processor(text=text, audio=[audio], max_length=131072,
                       audio_kwargs={"device": "cuda"}, return_tensors="pt").to("cuda")

with torch.inference_mode(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
    out = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=torch.ones_like(inputs["input_ids"]),
        input_features=inputs["input_features"],
        audio_feature_lengths=inputs["audio_feature_lengths"],
        audio_chunk_mapping=inputs["audio_chunk_mapping"],
        max_new_tokens=32768,
    )
transcript = processor.tokenizer.decode(
    out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
).strip()
# -> "[0.16][S01] Two two right, taxi via Bravo.[2.83][3.28][S02] ..."

4. The looping issue — and the recovery that handles it

Symptom. On near-silent stretches (quiet feeds, long gaps), greedy decoding occasionally enters a degenerate loop mid-transcript: a frozen [timestamp] repeated forever, a cycling token n-gram, or hallucinated micro-segments. Left alone, the decode burns its entire token budget on garbage and the rest of the clip is lost.

Handling (mirrors Whisper's temperature-fallback design). Detect the stall in-stream with a StoppingCriteria, truncate the output back to the last complete segment whose start timestamp still advances, then continue the same pass with sampling at the next temperature of a ladder — the continuation sees every previously emitted [Sxx] label, so speaker assignments stay consistent. Detection (pure Python):

import re, zlib

TS = re.compile(r"\[(\d+(?:\.\d+)?)\]")
SEGMENT = re.compile(r"\[(\d+(?:\.\d+)?)\]\[S\d+\].*?\[(\d+(?:\.\d+)?)\]", re.DOTALL)

def has_token_loop(generated, ngram=6, window=240, min_repeats=6):
    """Tail of the generated ids cycles: last `ngram` tokens recur >= `min_repeats` in the window."""
    if len(generated) < ngram * min_repeats:
        return False
    tail = list(generated[-window:]); probe = tail[-ngram:]
    return sum(tail[i:i+ngram] == probe for i in range(len(tail)-ngram, -1, -1)) >= min_repeats

def text_stalled(tail):
    """Frozen timestamp, or the text collapses under compression (ratio > 4)."""
    stamps = TS.findall(tail)
    if len(stamps) >= 6 and len(set(stamps[-6:])) == 1:
        return True
    raw = tail.encode()
    return len(raw) >= 600 and len(raw) / len(zlib.compress(raw)) > 4.0

def keep_token_count(pieces):
    """Tokens to keep after a stall: through the last segment whose start still advances."""
    text = "".join(pieces); keep_end = 0; prev = -1.0
    for m in SEGMENT.finditer(text):
        if float(m.group(1)) < prev:
            break
        prev = float(m.group(1)); keep_end = m.end()
    total = 0
    for i, p in enumerate(pieces):
        total += len(p)
        if total >= keep_end:
            return i + 1
    return len(pieces) if keep_end else 0

Recovery loop: run greedy with a stopping criterion that fires on has_token_loop(generated) or (every 25 steps) text_stalled(decoded_tail); on a stall, keep keep_token_count(...) tokens, append them to the prompt, and re-enter generation with do_sample=True at the next temperature of (0.2, 0.5, 0.7, 0.3, 0.6, 0.8), decrementing a shared max_new_tokens budget (stop when < 512 remains). In production this recovers a stalled 30-minute decode from ~45 emitted segments to full coverage in a single retry.

Residual behavior to detect downstream: on audio far outside the training distribution, diarization can go modal — every transmission minted as a new speaker, or everything merged into one. Both are trivially detectable (distinct-labels ≈ segment-count, or 1 label where context implies several) and a re-decode often lands on the correct mode.

5. Strongly recommended: a context-integration second step

VHF ATC audio is genuinely low quality — clipped AM radio, overlapping transmissions, weak receivers — so even a well-adapted acoustic model mishears callsigns and occasionally attributes a transmission to the wrong speaker. The single biggest quality jump comes from a second, non-acoustic step that reconciles the transcript against external context for the specific airport and frequency being processed:

• ADS-B traffic data (e.g. adsb.lol) for the clip's time window and location: the set of aircraft actually present is a closed candidate list — matching spoken callsigns against it names each speaker, corrects ASR-garbled callsigns, and merges/splits mismatched speaker labels using who could physically be on that frequency.
• The frequency's role (delivery / ground / tower / approach / departure): constrains the phraseology and which number groups are headings, altitudes, or frequencies rather than callsigns.
• Airport specifics (runway/taxiway identifiers, fixes) and airline telephony names (official designator-to-callword tables): repair garbled words toward what can actually be said on that frequency.

In practice this turns raw diarized output into a reliably attributed transcript — it is the step that fixes mismatched speakers, not more acoustic accuracy.