| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| from transformers import Gemma4Processor |
| from transformers.feature_extraction_utils import BatchFeature |
|
|
| from .parakeet_projector import PARAKEET_NAME, valid_frames_for |
|
|
| TARGET_SR = 16000 |
|
|
|
|
| def _load_audio(item, target_sr: int = TARGET_SR): |
| if isinstance(item, dict): |
| arr = np.asarray(item.get("array"), dtype="float32") |
| sr = int(item.get("sampling_rate") or target_sr) |
| elif isinstance(item, (str, bytes)): |
| import librosa |
| arr, sr = librosa.load(item, sr=target_sr, mono=True) |
| else: |
| arr, sr = np.asarray(item, dtype="float32"), target_sr |
| if arr.ndim > 1: |
| arr = arr.mean(axis=1) |
| if sr != target_sr: |
| import librosa |
| arr = librosa.resample(arr.astype("float32"), orig_sr=sr, |
| target_sr=target_sr) |
| return arr.astype("float32") |
|
|
|
|
| def _collect_audio(messages): |
| out = [] |
| for m in messages: |
| content = m.get("content") |
| if not isinstance(content, (list, tuple)): |
| continue |
| for part in content: |
| if isinstance(part, dict) and part.get("type") in ("audio", |
| "input_audio"): |
| out.append(part.get("audio", part.get("url"))) |
| return out |
|
|
|
|
| class LFG3Processor(Gemma4Processor): |
| parakeet_name = PARAKEET_NAME |
|
|
| def _feature_extractor(self): |
| if getattr(self, "_parakeet_fe", None) is None: |
| from transformers import AutoFeatureExtractor |
| self._parakeet_fe = AutoFeatureExtractor.from_pretrained( |
| self.parakeet_name |
| ) |
| return self._parakeet_fe |
|
|
| def _scaffold_ids(self): |
| tok = self.tokenizer |
| ids = tuple(tok.convert_tokens_to_ids(t) |
| for t in ("<|audio|>", "<|audio>", "<audio|>")) |
| if any(i is None or i == tok.unk_token_id for i in ids): |
| raise ValueError(f"audio scaffold tokens missing from vocab: {ids}") |
| return ids |
|
|
| def _batch_features(self, clips): |
| fe = self._feature_extractor() |
| per_clip, lengths = [], [] |
| for clip in clips: |
| out = fe(clip, sampling_rate=TARGET_SR, return_tensors="pt") |
| feats = out["input_features"][0] |
| mask = out.get("attention_mask") |
| n = int(mask[0].sum()) if mask is not None else feats.shape[0] |
| per_clip.append(feats[:n]) |
| lengths.append(n) |
|
|
| width = max(lengths) |
| padded = torch.zeros((len(per_clip), width, per_clip[0].shape[-1]), |
| dtype=per_clip[0].dtype) |
| mask = torch.zeros((len(per_clip), width), dtype=torch.long) |
| for i, (feats, n) in enumerate(zip(per_clip, lengths)): |
| padded[i, :n] = feats |
| mask[i, :n] = 1 |
| return {"input_features": padded, "encoder_attention_mask": mask} |
|
|
| def _splice(self, text: str, clip) -> tuple[list[int], int]: |
| ids = self.tokenizer(text, add_special_tokens=False)["input_ids"] |
| if clip is None: |
| return ids, 0 |
| marker_id, boa_id, eoa_id = self._scaffold_ids() |
| n = valid_frames_for(len(clip), TARGET_SR) |
| try: |
| pos = ids.index(marker_id) |
| except ValueError as exc: |
| raise ValueError( |
| "audio supplied but no <|audio|> marker in the rendered " |
| "prompt; include a {'type': 'audio'} part in the message" |
| ) from exc |
| return ids[:pos] + [boa_id] + [marker_id] * n + [eoa_id] + ids[pos + 1:], n |
|
|
| def apply_chat_template( |
| self, |
| conversations, |
| tokenize: bool = True, |
| return_dict: bool = True, |
| return_tensors: str = "pt", |
| padding: bool = False, |
| add_generation_prompt: bool = True, |
| enable_thinking: bool = False, |
| audio=None, |
| **kwargs, |
| ): |
| render = super().apply_chat_template |
| batched = bool(conversations) and isinstance(conversations[0], |
| (list, tuple)) |
| convs = list(conversations) if batched else [conversations] |
|
|
| texts = [render(c, tokenize=False, |
| add_generation_prompt=add_generation_prompt, |
| enable_thinking=enable_thinking, **kwargs) |
| for c in convs] |
| if not tokenize: |
| return texts if batched else texts[0] |
|
|
| rows, clips, frames = [], [], [] |
| for conv, text in zip(convs, texts): |
| found = _collect_audio(conv) if audio is None else list(audio) |
| if len(found) > 1: |
| raise ValueError("one audio clip per conversation is supported; " |
| f"got {len(found)}") |
| clip = _load_audio(found[0]) if found else None |
| ids, n = self._splice(text, clip) |
| rows.append(ids) |
| if clip is not None: |
| clips.append(clip) |
| frames.append(n) |
|
|
| pad_id = self.tokenizer.pad_token_id or 0 |
| width = max(len(r) for r in rows) |
| input_ids = [[pad_id] * (width - len(r)) + r for r in rows] |
| attn = [[0] * (width - len(r)) + [1] * len(r) for r in rows] |
|
|
| data = { |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), |
| "attention_mask": torch.tensor(attn, dtype=torch.long), |
| } |
| if clips: |
| if len(clips) != len(rows): |
| raise ValueError("every conversation in a batch must carry " |
| "audio, or none of them may") |
| data.update(self._batch_features(clips)) |
| data["valid_frames"] = torch.tensor(frames, dtype=torch.long) |
| return BatchFeature(data=data) |
|
|
