File size: 5,792 Bytes
7e5c0ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | 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)
|