remove mirrored code/ (use the GitHub repo)

code/data.py

@@ -1,302 +0,0 @@
-"""
-Data pipeline for the minimal SpeechLLM.
-
-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 `audio_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 downsampling arithmetic the model uses.
-
-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 audio_token_count(mel_len: int, n_downsample: int) -> int:
-    """mel frames -> number of LLM token slots this audio needs.
-
-    The connector concatenates every `n_downsample` adjacent frames and right-pads to
-    a multiple of that, so the count is ceil(frames / n_downsample), rounding up.
-    """
-    return -(-whisper_frames(mel_len) // n_downsample)
-
-
-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.
-    """
-
-    def __init__(self, manifest_filepaths, audio_root: str, ratios=None, seed: int = 0):
-        self.audio_root = audio_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.audio_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: int = 2
-    # 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) 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

@@ -1,201 +0,0 @@
-#!/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 \
-        --audio-root /path/to/data --limit 20
-
-"""
-
-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"))
-        # n_downsample used to be n_downsample_layers, an exponent: 1 meant 2 frames
-        # per token. Map the old key so earlier checkpoints still load.
-        n_downsample = hp.get("n_downsample") or 2 ** hp["n_downsample_layers"]
-        logger.info("  llm=%s encoder=%s n_downsample=%d frozen_encoder=%s",
-                    hp["llm_id"], hp["encoder_id"],
-                    n_downsample, 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=n_downsample, dtype=dtype,
-            # .get(): older checkpoints predate these options
-            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=n_downsample,
-            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)
-
-        # The trainable tensors (connector, LoRA) stay in float32 while the frozen encoder and
-        # LLM run in `dtype`, exactly as in training -- so the forward pass has to happen under
-        # autocast, or layer_norm sees a float32 weight and a float16 activation and raises
-        # "expected scalar type Half but found Float".
-        # passing inputs_embeds means `generate` returns only the new tokens --
-        # there is no prompt prefix to slice off
-        with torch.autocast(self.device.split(":")[0], dtype=self.dtype,
-                            enabled=not self.device.startswith("cpu")):
-            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("--audio-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.audio_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.audio_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

@@ -1,171 +0,0 @@
-"""
-A minimal SpeechLLM: Whisper encoder + connector + LLM with LoRA.
-
-    audio ──► Whisper encoder (trained or frozen) ──► connector (trained) ──┐
-                                                                            ├──► LLM + LoRA ──► text
-    text  ────────────────────────────────────────► embedding table ────────┘
-"""
-
-import logging
-
-import torch
-import torch.nn as nn
-from transformers import AutoModelForCausalLM, WhisperModel
-
-logger = logging.getLogger(__name__)
-
-
-class ConcatMLPConnector(nn.Module):
-    """Whisper hidden states → LLM-width embeddings by frame concatenation.
-    """
-
-    def __init__(self, d_encoder: int, d_llm: int, n_downsample: int = 2,
-                 hidden_dim: int = 0):
-        super().__init__()
-        self.n_downsample = n_downsample
-        hidden_dim = hidden_dim or d_llm
-        self.proj = nn.Sequential(
-            nn.LayerNorm(d_encoder * n_downsample),
-            nn.Linear(d_encoder * n_downsample, 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
-        k = self.n_downsample
-        pad = (-T) % k
-        if pad:
-            hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad))
-        stacked = hidden_states.reshape(B, (T + pad) // k, D * k)
-        return self.proj(stacked)                              # (B, ceil(T/k), d_llm)
-
-
-class SpeechLLM(nn.Module):
-    def __init__(self, llm_id: str, encoder_id: str,
-                 n_downsample: int = 2, dtype=torch.bfloat16,
-                 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
-
-        self.connector = ConcatMLPConnector(
-            d_encoder=self.encoder.config.d_model,
-            d_llm=self.llm.config.hidden_size,
-            n_downsample=n_downsample,
-            hidden_dim=adapter_hidden_dim,
-        )
-        logger.info("connector: %.1fM params",
-                    sum(p.numel() for p in self.connector.parameters()) / 1e6)
-
-        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",
-            )
-
-            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
-
-        if not freeze_encoder:
-            self.encoder.float()
-
-        if gradient_checkpointing:
-            ckpt_kwargs = {"use_reentrant": False}
-            self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=ckpt_kwargs)
-            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.
-        """
-        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()
-        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}