Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/1272-128104-0000.flac filter=lfs diff=lfs merge=lfs -text
+samples/1272-128104-0001.flac filter=lfs diff=lfs merge=lfs -text
+samples/1272-128104-0002.flac filter=lfs diff=lfs merge=lfs -text
+samples/1272-128104-0003.flac filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,45 @@
+---
+license: apache-2.0
+language: en
+tags: [speech-llm, desta, librispeech, tutorial]
+---
+
+# Interspeech tutorial — DeSTA-style SpeechLLM checkpoints
+
+Whisper-large-v3 encoder (frozen) → concat+MLP adapter → Qwen3-4B-Instruct-2507 + LoRA r32.
+Only the adapter and the LoRA weights are trained, so each checkpoint is ~147 MB; the base
+models are downloaded from their own repos at load time.
+
+| folder | training data | test-clean WER | test-clean gender |
+|---|---|---|---|
+| `selfgen_only` | 281k self-generated conversational responses (no task labels) | 3.93 (best prompt + post-processing) | 98.24 |
+| `selfgen_mix` | the same + 20% ASR + 5% gender SFT | 2.45 | 98.82 |
+
+The self-generation targets were written by Qwen3-4B given only the transcript and the
+speaker's gender: `<audio>{transcription} (Gender: {gender})</audio>` + "The audio is a
+passage read aloud from a book. Respond directly as a natural conversation partner. Do not
+mention the audio, the transcription, or the speaker attributes."
+
+`selfgen_only` has never seen a transcription instruction: it can transcribe, but needs an
+explicit prompt ("Transcribe the speech word for word. Output only the transcription …") and
+sometimes answers in the training-time metadata format. `selfgen_mix` follows the plain
+prompts directly.
+
+## Usage
+
+See the Colab notebook in the tutorial repo, or:
+
+```python
+from huggingface_hub import hf_hub_download, snapshot_download
+import sys, torch
+
+code = snapshot_download("kehanlu/interspeech-tutorial", allow_patterns="code/*")
+sys.path.insert(0, f"{code}/code")
+from inference import SpeechLLMForInference
+
+ckpt = hf_hub_download("kehanlu/interspeech-tutorial", "selfgen_mix/model.ckpt")
+pipe = SpeechLLMForInference.from_checkpoint(ckpt, dtype=torch.float16)   # float16 for a Colab T4
+print(pipe.generate([{"role": "user",
+                      "content": "<audio><|AUDIO|></audio>\n\nTranscribe the speech into text",
+                      "audios": [{"audio": "sample.flac"}]}]))
+```

code/data.py

@@ -0,0 +1,320 @@
+"""
+Data pipeline for the minimal SpeechLLM.
+
+The one idea to hold on to while reading this file: **the LLM never sees audio.**
+It sees a plain text token sequence in which every ``<|AUDIO|>`` has been replaced
+by exactly as many placeholder tokens as that audio will occupy after encoding.
+The model later overwrites those placeholders' *embeddings* with real speech
+features (see modeling.py). Nothing else about the LLM changes.
+
+Everything that turns "a manifest row" into "model inputs" happens in `Collator`.
+The `Dataset` is a dumb jsonl reader. Keeping it in one place matters here: the
+number of placeholders depends on the audio's true duration, which we only know
+once the waveform is loaded -- so the expansion has to happen next to the audio.
+
+Manifest format (one JSON object per line)::
+
+    {
+      "audios": [
+        {"audio_filepath": "ESC50/1-103999-A-30.wav"}
+      ],
+      "messages": [
+        {"role": "user", "content": "What do you hear? <|AUDIO|>"}
+      ],
+      "target": "A wooden door being knocked on."
+    }
+
+`audio_filepath` is relative to `data_root`. `target` may also be called
+`response` (both appear in the DeSTA3 manifests). Extra keys are ignored.
+"""
+
+import json
+import logging
+import math
+import os
+import random
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+import librosa
+import torch
+from torch.utils.data import Dataset
+
+logger = logging.getLogger(__name__)
+
+SAMPLE_RATE = 16000
+
+
+# --------------------------------------------------------------------------
+# How many tokens will one audio occupy?
+#
+# Whisper's feature extractor always pads to 30 s, so we cannot use the padded
+# length -- a 5 s clip would reserve 25 s worth of slots. We ask the feature
+# extractor for an attention mask over the mel frames instead, and push the true
+# length through the same conv arithmetic the model uses.
+#
+# These formulas mirror Qwen2AudioEncoder._get_feat_extract_output_lengths.
+# --------------------------------------------------------------------------
+
+def whisper_frames(mel_len: int) -> int:
+    """Whisper's encoder conv2 has stride 2. 3000 mel frames -> 1500 positions."""
+    return (mel_len - 1) // 2 + 1
+
+
+def conv1d_out(length: int, kernel: int = 3, stride: int = 2, padding: int = 1) -> int:
+    return (length + 2 * padding - kernel) // stride + 1
+
+
+def audio_token_count(mel_len: int, n_downsample_layers: int) -> int:
+    """mel frames -> number of LLM token slots this audio needs."""
+    length = whisper_frames(mel_len)
+    for _ in range(n_downsample_layers):
+        length = conv1d_out(length)
+    return length
+
+
+def expand_audio_locator(tokens: List[str], audio_locator: str,
+                         placeholder_token: str, slot_sizes: List[int]):
+    """Replace each `audio_locator` with `slot_sizes[i]` placeholder tokens.
+
+    ``["a", "<|AUDIO|>", "b"]`` with ``slot_sizes=[3]`` becomes
+    ``["a", "<|pad|>", "<|pad|>", "<|pad|>", "b"]`` and ``start_positions=[1]``.
+    """
+    out: List[str] = []
+    starts: List[int] = []
+    sizes = iter(slot_sizes)
+    for token in tokens:
+        if token == audio_locator:
+            starts.append(len(out))
+            out.extend([placeholder_token] * next(sizes))
+        else:
+            out.append(token)
+    return out, starts
+
+
+def resolve_manifest(path: str) -> str:
+    """Accept a local path or a huggingface.co/datasets/... resolve URL."""
+    prefix = "https://huggingface.co/datasets/"
+    if not path.startswith(prefix):
+        return path
+    from huggingface_hub import hf_hub_download
+
+    parts = path[len(prefix):].split("/")
+    repo_id, marker, revision = "/".join(parts[:2]), parts[2], parts[3]
+    assert marker == "resolve", f"unsupported HF url: {path}"
+    return hf_hub_download(repo_id=repo_id, filename="/".join(parts[4:]),
+                           revision=revision, repo_type="dataset")
+
+
+def resolve_audio_filepath(path: str) -> str:
+    if os.path.exists(path):
+        return path
+    wav = os.path.splitext(path)[0] + ".wav"
+    if os.path.exists(wav):
+        return wav
+    raise FileNotFoundError(path)
+
+
+def split_manifest_entries(entries):
+    """Manifest entries are a path (ratio 1.0) or {"path": ..., "ratio": ...}.
+
+    Returns (paths, ratios).
+    """
+    paths, ratios = [], []
+    for entry in entries:
+        if isinstance(entry, str):
+            path, ratio = entry, 1.0
+        else:
+            unknown = set(entry) - {"path", "ratio"}
+            assert "path" in entry and not unknown, f"bad manifest entry {entry!r}"
+            path, ratio = entry["path"], float(entry.get("ratio", 1.0))
+        assert ratio >= 0, f"{path}: ratio must be >= 0, got {ratio}"
+        paths.append(path)
+        ratios.append(ratio)
+    return paths, ratios
+
+
+def manifest_name(path: str) -> str:
+    return os.path.basename(path).removesuffix(".jsonl")
+
+
+class AudioTextDataset(Dataset):
+    """Lazy jsonl reader with per-manifest mixing ratios.
+
+    The DeSTA3 manifests run to hundreds of MB each, so we build a byte-offset
+    index up front and read one line at a time instead of holding the parsed
+    rows in memory.
+
+    `ratios[i]` scales manifest i: 0.3 keeps a random 30% of its rows, 2.5 keeps two
+    full copies plus a random 50%. The random part is drawn from (seed, epoch,
+    manifest) only -- never the rank -- so every DDP rank builds the same index
+    and DistributedSampler shards it consistently. Call `resample(epoch)` to draw a
+    fresh subset for a new epoch; the length does not change.
+    """
+
+    def __init__(self, manifest_filepaths, data_root: str, ratios=None, seed: int = 0):
+        self.data_root = data_root
+        self.paths = [resolve_manifest(p) for p in manifest_filepaths]
+        self.ratios = list(ratios) if ratios is not None else [1.0] * len(self.paths)
+        assert len(self.ratios) == len(self.paths)
+        self.seed = seed
+        self.rows = []           # per manifest: [(file_idx, byte_offset), ...]
+        for file_idx, path in enumerate(self.paths):
+            rows = []
+            with open(path, "rb") as f:
+                offset = 0
+                for line in f:
+                    if line.strip():
+                        rows.append((file_idx, offset))
+                    offset += len(line)
+            self.rows.append(rows)
+        self._handles = {}       # opened lazily, per dataloader worker
+        self.resample(epoch=0, log=True)
+
+    @property
+    def needs_resampling(self) -> bool:
+        """True when some ratio has a fractional part, i.e. a random subset is drawn."""
+        return any(r != math.floor(r) for r in self.ratios)
+
+    def resample(self, epoch: int, log: bool = False):
+        self.index = []          # (file_idx, byte_offset)
+        for file_idx, (rows, ratio) in enumerate(zip(self.rows, self.ratios)):
+            whole = math.floor(ratio)
+            n_extra = round((ratio - whole) * len(rows))
+            # str seed: hashed with sha512, identical across processes and runs
+            rng = random.Random(f"{self.seed}-{epoch}-{file_idx}")
+            self.index += rows * whole + rng.sample(rows, n_extra)
+            if log:
+                logger.info("  %-32s ratio %.2f: %d of %d rows", manifest_name(self.paths[file_idx]),
+                            ratio, whole * len(rows) + n_extra, len(rows))
+        if log:
+            logger.info("loaded %d rows from %d manifest(s)", len(self.index), len(self.paths))
+
+    def __len__(self):
+        return len(self.index)
+
+    def __getitem__(self, i):
+        file_idx, offset = self.index[i]
+        handle = self._handles.get(file_idx)
+        if handle is None:
+            handle = self._handles[file_idx] = open(self.paths[file_idx], "rb")
+        handle.seek(offset)
+        row = json.loads(handle.readline())
+
+        # DeSTA3 manifests use either key; `response` wins when both are present.
+        target = row.get("response") or row.get("target")
+        assert target, f"row {i} has neither `response` nor `target`"
+
+        audios = [{"audio_filepath": resolve_audio_filepath(
+                       os.path.join(self.data_root, audio["audio_filepath"]))}
+                  for audio in row["audios"]]
+        return {"messages": row["messages"], "target": target, "audios": audios}
+
+
+@dataclass
+class Collator:
+    """Turns a list of rows into the tensors `SpeechLLM.forward` expects.
+
+    Padding is on the **left** so that every sequence's answer ends at the same
+    index -- that keeps `generate()` simple. It also means every position we
+    record has to be shifted by that sequence's pad length.
+    """
+
+    tokenizer: Any
+    feature_extractor: Any
+    audio_locator: str = "<|AUDIO|>"
+    placeholder_token: str = "<|vision_pad|>"
+    max_seq_length: int = 1024
+    max_audio_seconds: float = 30.0
+    n_downsample_layers: int = 1
+    # inference: build the prompt only, with no answer appended and no labels.
+    # Same slot arithmetic as training -- that is the whole point of reusing this
+    # class rather than rebuilding the expansion in the demo script.
+    for_generation: bool = False
+
+    def __post_init__(self):
+        assert self.tokenizer.padding_side == "left"
+        assert self.placeholder_token in self.tokenizer.get_vocab(), (
+            f"placeholder_token {self.placeholder_token!r} is not in the tokenizer "
+            "vocabulary; pick a reserved/unused token of your LLM")
+        self.placeholder_id = self.tokenizer.convert_tokens_to_ids(self.placeholder_token)
+        self.pad_id = self.tokenizer.pad_token_id
+
+    def __call__(self, batch):
+        # 1. audio -> mel features, and the true (unpadded) length of each
+        waveforms = []
+        for row in batch:
+            for audio in row["audios"]:
+                wav, _ = librosa.load(audio["audio_filepath"], sr=SAMPLE_RATE, mono=True)
+                waveforms.append(wav[:int(self.max_audio_seconds * SAMPLE_RATE)])
+
+        features = self.feature_extractor(
+            waveforms, sampling_rate=SAMPLE_RATE,
+            return_tensors="pt", return_attention_mask=True,
+        )
+        mel_lengths = features["attention_mask"].sum(-1).tolist()
+        audio_lengths = [audio_token_count(n, self.n_downsample_layers) for n in mel_lengths]
+
+        # 2. text -> token ids, with <|AUDIO|> expanded to placeholders
+        context_ids, target_ids = [], []
+        start_positions, audio_index = [], 0
+
+        for row in batch:
+            n_audios = len(row["audios"])
+            # each audio reserves one slot per speech feature frame
+            slot_sizes = audio_lengths[audio_index:audio_index + n_audios]
+
+            prompt = self.tokenizer.apply_chat_template(
+                row["messages"], tokenize=False, add_generation_prompt=True,
+                enable_thinking=False,
+            )
+            tokens = self.tokenizer.tokenize(prompt)
+            assert tokens.count(self.audio_locator) == n_audios, (
+                f"{n_audios} audios but {tokens.count(self.audio_locator)} "
+                f"{self.audio_locator} in the prompt")
+
+            tokens, starts = expand_audio_locator(
+                tokens, self.audio_locator, self.placeholder_token, slot_sizes)
+
+            # straight to ids -- no convert_tokens_to_string round-trip, so the
+            # placeholder can never be re-tokenized into something else
+            ctx = self.tokenizer.convert_tokens_to_ids(tokens)
+            if self.for_generation:
+                tgt = []                       # nothing to condition on; the model writes it
+            else:
+                tgt = self.tokenizer.encode(row["target"], add_special_tokens=False)
+                tgt = tgt + [self.tokenizer.eos_token_id]
+
+            assert len(ctx) < self.max_seq_length, (
+                f"prompt alone is {len(ctx)} tokens (max_seq_length="
+                f"{self.max_seq_length}); shorten the audio or raise the limit")
+            tgt = tgt[:self.max_seq_length - len(ctx)]      # only ever truncate the answer
+
+            context_ids.append(ctx)
+            target_ids.append(tgt)
+            start_positions.append(starts)
+            audio_index += n_audios
+
+        # 3. left-pad, build labels, shift the recorded positions
+        width = max(len(c) + len(t) for c, t in zip(context_ids, target_ids))
+        input_ids = torch.full((len(batch), width), self.pad_id, dtype=torch.long)
+        attention_mask = torch.zeros((len(batch), width), dtype=torch.long)
+        labels = torch.full((len(batch), width), -100, dtype=torch.long)
+        shifted_starts = []
+
+        for i, (ctx, tgt) in enumerate(zip(context_ids, target_ids)):
+            pad = width - len(ctx) - len(tgt)
+            input_ids[i, pad:] = torch.tensor(ctx + tgt, dtype=torch.long)
+            attention_mask[i, pad:] = 1
+            labels[i, pad + len(ctx):] = torch.tensor(tgt, dtype=torch.long)  # answer only
+            for start in start_positions[i]:
+                shifted_starts.append((i, start + pad))
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "input_features": features["input_features"],
+            "audio_lengths": audio_lengths,          # per audio, after downsampling
+            "start_positions": shifted_starts,       # per audio, (row, first placeholder)
+        }

code/inference.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python
+"""
+Inference for the minimal SpeechLLM -- the tutorial's demo entry point.
+
+    # one file
+    python inference.py --ckpt exp/run/checkpoints/last.ckpt --audio sample.flac
+
+    # score a manifest: prints reference vs hypothesis side by side
+    python inference.py --ckpt exp/run/checkpoints/last.ckpt \
+        --manifest data/Librispeech-dev-test/test-clean_asr.jsonl \
+        --data-root /home/u8915687/work/data/audios2 --limit 20
+
+The architecture is read back out of the checkpoint's `hyper_parameters`, never
+hand-written here. That matters more than it looks: `SpeechLLMModule` loads with
+`strict=False`, and the connector's downsampling stack is an `nn.Sequential`
+whose keys are positional (`downsample.0.*`, `downsample.2.*`). Build the model
+with the wrong `n_downsample_layers` and those tensors become "missing keys",
+which `strict=False` silently drops -- you get a randomly-initialised connector,
+no error, and fluent nonsense out. `_load_trainable` below turns that into a
+hard failure instead.
+
+Unlike DeSTA3 there is no VAD and no ASR decoder in the loop: DeSTA3 runs Whisper's
+decoder to produce a transcript, and needs VAD to decide which clips are speech
+at all. Here the LLM only sees the speech features, exactly as in training.
+"""
+
+import argparse
+import json
+import logging
+import os
+
+import torch
+from transformers import AutoFeatureExtractor, AutoTokenizer
+
+from data import Collator
+from modeling import SpeechLLM
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+DEFAULT_PROMPT = "<audio><|AUDIO|></audio>\n\nTranscribe the speech into text"
+
+
+class SpeechLLMForInference:
+    """Checkpoint in, text out.
+
+    `generate()` takes DeSTA3-style messages so the demo reads like a chat call:
+
+        [{"role": "user",
+          "content": "<audio><|AUDIO|></audio>\\n\\nTranscribe the speech into text",
+          "audios": [{"audio": "sample.flac"}]}]
+
+    Pass a list of those to batch several utterances in one forward pass.
+    """
+
+    def __init__(self, model, tokenizer, collator, device, dtype):
+        self.model, self.tokenizer = model, tokenizer
+        self.collator, self.device, self.dtype = collator, device, dtype
+
+    # -- loading ----------------------------------------------------------
+    @classmethod
+    def from_checkpoint(cls, ckpt_path, device=None, dtype=None):
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        hp = ckpt["hyper_parameters"]
+        dtype = dtype or getattr(torch, hp["dtype"])
+        logger.info("checkpoint %s (epoch %s, step %s)",
+                    os.path.basename(ckpt_path), ckpt.get("epoch"), ckpt.get("global_step"))
+        logger.info("  llm=%s encoder=%s adapter=%s n_downsample=%d frozen_encoder=%s",
+                    hp["llm_id"], hp["encoder_id"], hp.get("adapter_type", "cnn"),
+                    hp["n_downsample_layers"], hp.get("freeze_encoder", False))
+
+        # rebuilt exactly as SpeechLLMModule.__init__ does, or the ids shift
+        tokenizer = AutoTokenizer.from_pretrained(hp["llm_id"])
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.padding_side = "left"
+        tokenizer.add_tokens([hp["audio_locator"]])
+        feature_extractor = AutoFeatureExtractor.from_pretrained(hp["encoder_id"])
+
+        model = SpeechLLM(
+            llm_id=hp["llm_id"], encoder_id=hp["encoder_id"],
+            n_downsample_layers=hp["n_downsample_layers"], dtype=dtype,
+            # .get(): checkpoints from before these options were cnn + trained encoder
+            adapter_type=hp.get("adapter_type", "cnn"),
+            adapter_hidden_dim=hp.get("adapter_hidden_dim", 0),
+            freeze_encoder=hp.get("freeze_encoder", False), use_lora=True,
+            lora_rank=hp["lora_rank"], lora_alpha=hp["lora_alpha"],
+            lora_dropout=hp["lora_dropout"],
+            lora_target_modules=hp["lora_target_modules"],
+        )
+        cls._load_trainable(model, ckpt["state_dict"])
+        model.to(device).eval()
+
+        collator = Collator(
+            tokenizer=tokenizer, feature_extractor=feature_extractor,
+            audio_locator=hp["audio_locator"], placeholder_token=hp["placeholder_token"],
+            max_seq_length=hp["max_seq_length"], max_audio_seconds=hp["max_audio_seconds"],
+            n_downsample_layers=hp["n_downsample_layers"],
+            for_generation=True,
+        )
+        return cls(model, tokenizer, collator, device, dtype)
+
+    @staticmethod
+    def _load_trainable(model, state_dict):
+        """Load the trainable-only checkpoint, and prove every tensor landed."""
+        state = {k[len("model."):]: v for k, v in state_dict.items() if k.startswith("model.")}
+        expected = {n for n, p in model.named_parameters() if p.requires_grad}
+        missing, unexpected = expected - set(state), set(state) - expected
+        assert not unexpected, (
+            f"{len(unexpected)} tensors in the checkpoint match nothing in the model, "
+            f"e.g. {sorted(unexpected)[:3]} -- the architecture does not match")
+        assert not missing, (
+            f"{len(missing)} trainable tensors were not in the checkpoint, "
+            f"e.g. {sorted(missing)[:3]} -- they would stay randomly initialised")
+        result = model.load_state_dict(state, strict=False)
+        assert not result.unexpected_keys, result.unexpected_keys
+        logger.info("  restored %d trainable tensors (%.1fM params)", len(state),
+                    sum(v.numel() for v in state.values()) / 1e6)
+
+    # -- generation -------------------------------------------------------
+    def _to_rows(self, conversations):
+        """DeSTA3-style messages -> the row dicts `Collator` consumes."""
+        rows = []
+        for conv in conversations:
+            audios = []
+            for message in conv:
+                for audio in message.get("audios", []):
+                    path = audio["audio"] if isinstance(audio, dict) else audio
+                    assert os.path.exists(path), f"no such audio: {path}"
+                    audios.append({"audio_filepath": path})
+            n_locators = sum(m["content"].count(self.collator.audio_locator) for m in conv)
+            assert n_locators == len(audios), (
+                f"{len(audios)} audios but {n_locators} {self.collator.audio_locator} "
+                "in the conversation")
+            # strip `audios` before the chat template ever sees it
+            rows.append({"messages": [{"role": m["role"], "content": m["content"]} for m in conv],
+                         "audios": audios, "target": ""})
+        return rows
+
+    @torch.no_grad()
+    def generate(self, conversations, max_new_tokens=200, do_sample=False, **generation_kwargs):
+        if conversations and isinstance(conversations[0], dict):
+            conversations = [conversations]          # a single conversation
+        batch = self.collator(self._to_rows(conversations))
+
+        batch["input_ids"] = batch["input_ids"].to(self.device)
+        batch["attention_mask"] = batch["attention_mask"].to(self.device)
+        # SpeechLLM.encode_audio casts features to the encoder's own dtype
+        batch["input_features"] = batch["input_features"].to(self.device)
+
+        # passing inputs_embeds means `generate` returns only the new tokens --
+        # there is no prompt prefix to slice off
+        generated = self.model.generate(
+            batch, self.tokenizer, max_new_tokens=max_new_tokens,
+            do_sample=do_sample, **generation_kwargs)
+        return [t.strip() for t in
+                self.tokenizer.batch_decode(generated, skip_special_tokens=True)]
+
+
+# ------------------------------------------------------------------ main ---
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--audio", nargs="+", help="one or more audio files")
+    ap.add_argument("--manifest", help="jsonl to run over instead of --audio")
+    ap.add_argument("--data-root", default="")
+    ap.add_argument("--limit", type=int, default=10, help="rows to take from --manifest")
+    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
+    ap.add_argument("--batch-size", type=int, default=4)
+    ap.add_argument("--max-new-tokens", type=int, default=200)
+    ap.add_argument("--device", default=None)
+    args = ap.parse_args()
+    assert args.audio or args.manifest, "give --audio or --manifest"
+
+    pipe = SpeechLLMForInference.from_checkpoint(args.ckpt, device=args.device)
+
+    if args.audio:
+        items = [(os.path.join(args.data_root, p), None) for p in args.audio]
+    else:
+        items = []
+        with open(args.manifest) as f:
+            for line in f:
+                if len(items) >= args.limit:
+                    break
+                row = json.loads(line)
+                items.append((os.path.join(args.data_root,
+                                           row["audios"][0]["audio_filepath"]),
+                              row.get("response") or row.get("target")))
+
+    for start in range(0, len(items), args.batch_size):
+        chunk = items[start:start + args.batch_size]
+        convs = [[{"role": "user", "content": args.prompt,
+                   "audios": [{"audio": path}]}] for path, _ in chunk]
+        for (path, reference), hypothesis in zip(chunk, pipe.generate(
+                convs, max_new_tokens=args.max_new_tokens)):
+            print(f"\n--- {os.path.basename(path)}")
+            if reference is not None:
+                print(f"  ref: {reference}")
+            print(f"  hyp: {hypothesis}")
+
+
+if __name__ == "__main__":
+    main()

code/modeling.py

@@ -0,0 +1,252 @@
+"""
+A minimal SpeechLLM: Whisper encoder + connector + LLM with LoRA.
+
+    audio ──► Whisper encoder (trained or frozen) ──► connector (trained) ──┐
+                                                                            ├──► LLM + LoRA ──► text
+    text  ────────────────────────────────────────► embedding table ────────┘
+
+Two connectors (`adapter_type`), both 50 fps -> 25 fps per downsampling step:
+    cnn         strided Conv1d + GELU, then LayerNorm + Linear
+    concat_mlp  concatenate adjacent frames, then LayerNorm + Linear + GELU + Linear
+
+A frozen LLM does *not* make training cheap: the speech features enter at layer 0,
+so gradients still flow back through every LLM layer to reach the connector, and
+every LLM activation has to be kept for the backward pass. Turn on
+`gradient_checkpointing` if that hurts.
+"""
+
+import logging
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, WhisperModel
+
+logger = logging.getLogger(__name__)
+
+
+class SpeechConnector(nn.Module):
+    """Whisper hidden states → LLM-width embeddings.
+
+    Takes the encoder's last hidden state and applies two steps:
+
+    1. **Strided Conv1d downsampling.** Whisper emits 50 frames/s, which is far
+       too many tokens to hand an LLM. One stride-2 convolution brings that to
+       25 frames/s -- 40 ms per token, the same rate as Qwen2-Audio. Each extra
+       layer halves the rate again.
+
+    2. **Linear projection to the LLM's width** (from Qwen2-Audio, whose entire
+       connector is ``AvgPool1d(2)`` + ``Linear``).
+    """
+
+    def __init__(self, d_encoder: int, d_llm: int, n_downsample_layers: int = 1):
+        super().__init__()
+        layers = []
+        for _ in range(n_downsample_layers):
+            layers += [nn.Conv1d(d_encoder, d_encoder, kernel_size=3, stride=2, padding=1),
+                       nn.GELU()]
+        self.downsample = nn.Sequential(*layers)
+
+        self.proj = nn.Sequential(nn.LayerNorm(d_encoder), nn.Linear(d_encoder, d_llm))
+
+    def forward(self, hidden_states):                          # (B, T, d_encoder)
+        hidden_states = self.downsample(hidden_states.transpose(1, 2)).transpose(1, 2)
+        return self.proj(hidden_states)                        # (B, T', d_llm)
+
+
+class ConcatMLPConnector(nn.Module):
+    """Whisper hidden states → LLM-width embeddings by frame concatenation.
+
+    Every `2 ** n_downsample_layers` adjacent frames are concatenated into one
+    feature (d_encoder * k), which a two-layer MLP maps to the LLM's width. This is
+    the SLAM-ASR / LLaVA-style projector: no convolution, all mixing happens in the
+    MLP. The sequence is right-padded to a multiple of k, so a length-L input gives
+    ceil(L / k) tokens -- the same count as the CNN, so `data.audio_token_count`
+    serves both.
+    """
+
+    def __init__(self, d_encoder: int, d_llm: int, n_downsample_layers: int = 1,
+                 hidden_dim: int = 0):
+        super().__init__()
+        self.k = 2 ** n_downsample_layers
+        hidden_dim = hidden_dim or d_llm
+        self.proj = nn.Sequential(
+            nn.LayerNorm(d_encoder * self.k),
+            nn.Linear(d_encoder * self.k, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, d_llm),
+        )
+
+    def forward(self, hidden_states):                          # (B, T, d_encoder)
+        B, T, D = hidden_states.shape
+        pad = (-T) % self.k
+        if pad:
+            hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad))
+        stacked = hidden_states.reshape(B, (T + pad) // self.k, D * self.k)
+        return self.proj(stacked)                              # (B, ceil(T/k), d_llm)
+
+
+ADAPTERS = {"cnn": SpeechConnector, "concat_mlp": ConcatMLPConnector}
+
+
+class SpeechLLM(nn.Module):
+    def __init__(self, llm_id: str, encoder_id: str,
+                 n_downsample_layers: int = 1, dtype=torch.bfloat16,
+                 adapter_type: str = "cnn", adapter_hidden_dim: int = 0,
+                 gradient_checkpointing: bool = False, cache_dir=None,
+                 freeze_encoder: bool = True, use_lora: bool = False,
+                 lora_rank: int = 32, lora_alpha: int = 64, lora_dropout: float = 0.05,
+                 lora_target_modules=("q_proj", "k_proj", "v_proj", "o_proj")):
+        super().__init__()
+        self.llm = AutoModelForCausalLM.from_pretrained(
+            llm_id, torch_dtype=dtype, cache_dir=cache_dir)
+        # We only ever need the encoder; loading WhisperModel and keeping .encoder
+        # lets the decoder weights fall out of scope immediately.
+        self.encoder = WhisperModel.from_pretrained(
+            encoder_id, torch_dtype=dtype, cache_dir=cache_dir).encoder
+
+        assert adapter_type in ADAPTERS, f"adapter_type must be one of {list(ADAPTERS)}"
+        extra = {"hidden_dim": adapter_hidden_dim} if adapter_type == "concat_mlp" else {}
+        self.connector = ADAPTERS[adapter_type](
+            d_encoder=self.encoder.config.d_model,
+            d_llm=self.llm.config.hidden_size,
+            n_downsample_layers=n_downsample_layers,
+            **extra,
+        )
+        logger.info("adapter %s: %.1fM params", adapter_type,
+                    sum(p.numel() for p in self.connector.parameters()) / 1e6)
+
+        # --- what trains ---------------------------------------------------
+        # The connector always does. The LLM trains only through LoRA adapters,
+        # if at all. The encoder is a switch: leaving it frozen is much cheaper
+        # than it looks, because with every encoder parameter frozen and
+        # `input_features` not requiring grad, PyTorch builds no graph for the
+        # encoder at all -- its activations are never stored. Unfreezing it
+        # turns 32 Whisper layers back into stored activations, which costs far
+        # more memory than the 635M parameters themselves.
+        for param in self.llm.parameters():
+            param.requires_grad = False
+
+        self.use_lora = use_lora
+        if use_lora:
+            from peft import LoraConfig ,get_peft_model
+            lora_config = LoraConfig(
+                r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
+                target_modules=list(lora_target_modules), bias="none",
+                task_type="CAUSAL_LM",
+            )
+            # unwrap back to a plain module: we call the LLM directly (with
+            # inputs_embeds) rather than through the PeftModel wrapper
+            self.llm = get_peft_model(self.llm, lora_config).base_model.model
+            n_lora = sum(p.numel() for n, p in self.llm.named_parameters() if "lora_" in n)
+            assert n_lora > 0, f"LoRA matched no modules in {lora_target_modules}"
+            logger.info("LoRA r=%d on %s: %.1fM adapter params",
+                        lora_rank, list(lora_target_modules), n_lora / 1e6)
+
+        self.freeze_encoder = freeze_encoder
+        for param in self.encoder.parameters():
+            param.requires_grad = not freeze_encoder
+
+        # Standard mixed precision: frozen weights stay in `dtype`, everything that
+        # trains is stored in float32 and autocast does the bf16 compute. bf16 has
+        # only 7 fraction bits (~0.8% between neighbouring values), so AdamW's
+        # small per-step updates would round to zero on bf16 weights.
+        # The connector and LoRA are created in float32 already; the encoder is
+        # loaded in `dtype`, so cast it when it trains.
+        if not freeze_encoder:
+            self.encoder.float()
+
+        if gradient_checkpointing:
+            # use_reentrant=False matters under DDP: the reentrant implementation
+            # runs the recomputation outside autograd's view, so DDP can mark a
+            # gradient ready twice and abort with "Expected to mark a variable
+            # ready only once". The non-reentrant path composes with DDP.
+            ckpt_kwargs = {"use_reentrant": False}
+            self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=ckpt_kwargs)
+            # deliberately no enable_input_require_grads(): the spliced speech
+            # features already make inputs_embeds require grad, and that hook
+            # would make it a leaf, breaking the splice (see build_inputs_embeds)
+            if not freeze_encoder:
+                self.encoder.gradient_checkpointing_enable(
+                    gradient_checkpointing_kwargs=ckpt_kwargs)
+
+        dtypes = {}
+        for name, param in self.named_parameters():
+            if param.requires_grad:
+                part = name.split(".")[0]
+                dtypes.setdefault(part, set()).add(str(param.dtype).replace("torch.", ""))
+        logger.info("trainable dtypes: %s", {k: sorted(v) for k, v in sorted(dtypes.items())})
+
+        by_part = {}
+        for name, param in self.named_parameters():
+            if param.requires_grad:
+                by_part[name.split(".")[0]] = by_part.get(name.split(".")[0], 0) + param.numel()
+        trainable = sum(by_part.values())
+        total = sum(p.numel() for p in self.parameters())
+        logger.info("trainable %.1fM / %.1fM total (%.2f%%) -- %s",
+                    trainable / 1e6, total / 1e6, 100 * trainable / total,
+                    {k: f"{v/1e6:.1f}M" for k, v in sorted(by_part.items())})
+
+    def encode_audio(self, input_features):
+        """mel features → speech embeddings in the LLM's width.
+
+        We run the encoder over the full 30 s window including padding, and crop
+        afterwards. Whisper is *trained* that way, so this is its native regime;
+        masking the padding would be off-distribution for a frozen encoder.
+        """
+        input_features = input_features.to(self.encoder.conv1.weight.dtype)
+        return self.connector(self.encoder(input_features).last_hidden_state)
+
+    def build_inputs_embeds(self, input_ids, speech_features, audio_lengths, start_positions):
+        """Overwrite each audio's placeholder embeddings with its speech features."""
+        embed = self.llm.get_input_embeddings()
+        # The embedding table is frozen, so `embed(input_ids)` comes back as a
+        # leaf with requires_grad=False -- and writing in-place into a leaf that
+        # requires grad is illegal, which is exactly what happens once gradient
+        # checkpointing turns that flag on. Clone to get a non-leaf we may write
+        # into. Splicing the speech features in is what makes the result require
+        # grad, which is also what checkpointing needs from the LLM's input.
+        inputs_embeds = embed(input_ids).clone()
+
+        for k, (row, start) in enumerate(start_positions):
+            segment = speech_features[k, :audio_lengths[k]]
+            end = start + segment.size(0)
+            assert end <= inputs_embeds.size(1), (
+                f"audio {k} needs slots [{start}:{end}] but the sequence is "
+                f"{inputs_embeds.size(1)} long")
+            # outside autocast (inference) the float32 connector output must be
+            # cast to the LLM's embedding dtype before the in-place write
+            inputs_embeds[row, start:end] = segment.to(inputs_embeds.dtype)
+
+        return inputs_embeds
+
+    def forward(self, input_ids, attention_mask, input_features, audio_lengths,
+                start_positions, labels=None):
+        speech_features = self.encode_audio(input_features)
+        inputs_embeds = self.build_inputs_embeds(
+            input_ids, speech_features, audio_lengths, start_positions)
+        return self.llm(inputs_embeds=inputs_embeds,
+                        attention_mask=attention_mask,
+                        labels=labels)
+
+    @torch.no_grad()
+    def generate(self, batch, tokenizer, **generation_kwargs):
+        """Greedy/sampled decoding from the same batch dict used for training.
+
+        Note we pass `inputs_embeds`, so `generate` returns only the newly
+        generated ids -- there is no prompt prefix to strip.
+        """
+        speech_features = self.encode_audio(batch["input_features"])
+        inputs_embeds = self.build_inputs_embeds(
+            batch["input_ids"], speech_features, batch["audio_lengths"],
+            batch["start_positions"])
+        return self.llm.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=batch["attention_mask"],
+            pad_token_id=tokenizer.pad_token_id,
+            **generation_kwargs,
+        )
+
+    def trainable_state_dict(self):
+        """Only the connector -- a few tens of MB instead of the full ~11 GB."""
+        return {name: param.detach().clone()
+                for name, param in self.named_parameters() if param.requires_grad}

samples/1272-128104-0000.flac

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e25e22555cd16e90edb0a3b49fdcf1fe652b2a1250ab643634db33895c75b41
+size 120041

samples/1272-128104-0001.flac

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46a9d58622b4675b29564da2d9ba73e702241c5fa969f12c387cad4aa984276a
+size 101672

samples/1272-128104-0002.flac

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59bf7bb1129c11734a8dd362ad52b4e10ae31c03f0879fc92f7107ecd513513f
+size 255853

samples/1272-128104-0003.flac

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60d72657439cd3092fc2313efe0fabd31f0cd2f28f8ca46a21380c290b2e031d
+size 196122

samples/samples.json

@@ -0,0 +1,18 @@
+[
+  {
+    "file": "1272-128104-0000.flac",
+    "transcription": "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel"
+  },
+  {
+    "file": "1272-128104-0001.flac",
+    "transcription": "nor is mister quilter's manner less interesting than his matter"
+  },
+  {
+    "file": "1272-128104-0002.flac",
+    "transcription": "he tells us that at this festive season of the year with christmas and roast beef looming before us similes drawn from eating and its results occur most readily to the mind"
+  },
+  {
+    "file": "1272-128104-0003.flac",
+    "transcription": "he has grave doubts whether sir frederick leighton's work is really greek after all and can discover in it but little of rocky ithaca"
+  }
+]

selfgen_mix/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a40967ad40705a5fb5bc08d8d50bc8eaa5a6ef05f2915a15e68c5d60428fa4a
+size 146925138

selfgen_only/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbea96fded8610bd5d66b5576c3d8c5635b2ba3e063896bd32c9ec562910c548
+size 146924882