File size: 3,851 Bytes
42c890c 68b4041 42c890c 2643de9 42c890c | 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 | from typing import List, Optional, Union
import numpy as np
from transformers import AutoProcessor, AutoTokenizer, WhisperFeatureExtractor
from transformers.feature_extraction_utils import BatchFeature
AUDIO_TOKENS_PER_SEGMENT = 750 # must match the projector output length in modeling_yor_audio.py
AUDIO_SPECIAL_TOKENS = ["<|audio_bos|>", "<|audio_eos|>", "<|AUDIO|>"]
SAMPLE_RATE = 16_000
NUM_MEL_BINS = 128
HOP_LENGTH = 160
CHUNK_LENGTH = 30
class YorAudioProcessor:
audio_token = "<|AUDIO|>"
audio_bos_token = "<|audio_bos|>"
audio_eos_token = "<|audio_eos|>"
def __init__(
self,
afriqueqwen_path: str = "McGill-NLP/AfriqueQwen3.5-9B-50Langs-Instruct-v1",
feature_extractor: Optional[WhisperFeatureExtractor] = None,
):
self.feature_extractor = feature_extractor or WhisperFeatureExtractor(
feature_size=NUM_MEL_BINS,
sampling_rate=SAMPLE_RATE,
hop_length=HOP_LENGTH,
chunk_length=CHUNK_LENGTH,
n_fft=400,
)
self.tokenizer = AutoTokenizer.from_pretrained(afriqueqwen_path)
self._add_audio_tokens()
def _add_audio_tokens(self):
missing = [t for t in AUDIO_SPECIAL_TOKENS if t not in self.tokenizer.get_vocab()]
if missing:
self.tokenizer.add_special_tokens({"additional_special_tokens": missing})
@property
def audio_token_id(self) -> int:
return self.tokenizer.convert_tokens_to_ids(self.audio_token)
def build_prompt(self, text: str, num_audios: int = 1) -> str:
audio_block = (
self.audio_bos_token
+ self.audio_token * AUDIO_TOKENS_PER_SEGMENT
+ self.audio_eos_token
)
return audio_block * num_audios + text
def __call__(
self,
text: Union[str, List[str]],
audio: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
sampling_rate: int = SAMPLE_RATE,
return_tensors: str = "pt",
padding: bool = True,
truncation: bool = True,
max_length: Optional[int] = None,
) -> BatchFeature:
if isinstance(text, str):
text = [text]
encoding = self.tokenizer(
text,
return_tensors=return_tensors,
padding=padding,
truncation=truncation,
max_length=max_length,
)
result = BatchFeature(data=dict(encoding))
if audio is not None:
if isinstance(audio, np.ndarray):
audio = [audio]
if audio[0].ndim == 2:
import torch
result["input_features"] = torch.tensor(np.stack(audio))
else:
audio_out = self.feature_extractor(
audio, sampling_rate=sampling_rate, return_tensors=return_tensors
)
result["input_features"] = audio_out["input_features"]
return result
def decode(self, token_ids, **kwargs) -> str:
return self.tokenizer.decode(token_ids, **kwargs)
def batch_decode(self, sequences, **kwargs) -> List[str]:
return self.tokenizer.batch_decode(sequences, **kwargs)
def save_pretrained(self, save_directory: str):
self.feature_extractor.save_pretrained(save_directory)
self.tokenizer.save_pretrained(save_directory)
@classmethod
def from_pretrained(cls, path: str, **kwargs) -> "YorAudioProcessor":
processor = cls.__new__(cls)
processor.feature_extractor = WhisperFeatureExtractor.from_pretrained(path)
processor.tokenizer = AutoTokenizer.from_pretrained(path)
processor._add_audio_tokens()
return processor
@classmethod
def register_for_auto_class(cls, auto_class="AutoProcessor"):
pass
AutoProcessor.register("yor_audio", YorAudioProcessor)
|