| 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 |
| 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) |
|
|
