File size: 2,767 Bytes
72ecc83 | 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 | """
Wren-ASR processor: audio → Mimi codes (and optionally back to text via the
tokenizer for decoding model outputs).
Usage:
processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
inputs = processor(audio=wav, sampling_rate=sr) # → {"audio_codes": [k, T]}
ids = model.generate(**inputs, max_new_tokens=200)
text = processor.batch_decode(ids, skip_special_tokens=True)[0]
"""
from typing import Optional, Union
import torch
from transformers.processing_utils import ProcessorMixin
class WrenASRProcessor(ProcessorMixin):
attributes = ["tokenizer"]
tokenizer_class = "AutoTokenizer"
def __init__(self, tokenizer, mimi_model_name: str = "kyutai/mimi", k_codebooks: int = 3, **kwargs):
super().__init__(tokenizer=tokenizer)
self.mimi_model_name = mimi_model_name
self.k_codebooks = k_codebooks
self._mimi = None
@property
def mimi(self):
if self._mimi is None:
from transformers import MimiModel
self._mimi = MimiModel.from_pretrained(self.mimi_model_name).eval()
for p in self._mimi.parameters():
p.requires_grad_(False)
return self._mimi
@torch.no_grad()
def __call__(
self,
audio: Optional[torch.Tensor] = None,
sampling_rate: Optional[int] = None,
audio_codes: Optional[torch.LongTensor] = None,
**kwargs,
):
"""Either pass `audio` (raw waveform) + `sampling_rate`, or pre-computed
`audio_codes` of shape [k, T] / [B, k, T].
Returns: {"audio_codes": LongTensor [B, k, T]}.
"""
if audio_codes is not None:
codes = audio_codes
if codes.dim() == 2:
codes = codes.unsqueeze(0)
return {"audio_codes": codes}
if audio is None:
raise ValueError("Provide either `audio` (waveform) or `audio_codes`.")
if sampling_rate is None:
raise ValueError("`sampling_rate` is required when passing `audio`.")
wav = audio
if wav.dim() == 1:
wav = wav.unsqueeze(0)
if sampling_rate != 24000:
import torchaudio.transforms as T
wav = T.Resample(sampling_rate, 24000)(wav)
x = wav.unsqueeze(0) # [1, 1, T]
out = self.mimi.encode(x, num_quantizers=self.k_codebooks)
codes = out.audio_codes # [1, k, T]
return {"audio_codes": codes}
def batch_decode(self, *args, **kwargs):
return self.tokenizer.batch_decode(*args, **kwargs)
def decode(self, *args, **kwargs):
return self.tokenizer.decode(*args, **kwargs)
|