| """AprielVoiceProcessor — Whisper audio, Gemma4 images, and text processing.""" |
|
|
| import math |
| import re |
| from typing import Any, Dict, List, Optional, Tuple, Union |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import transformers |
|
|
| from .apriel_voice_config import AprielVoiceConfig |
|
|
| _SHORT_AUDIO_THRESHOLD_SECONDS = 2.0 |
| _SHORT_AUDIO_SILENCE_SECONDS = 3.0 |
|
|
|
|
| def _prepend_silence_to_short_audios( |
| audios: List[np.ndarray], sampling_rate: int |
| ) -> List[np.ndarray]: |
| threshold_samples = int(_SHORT_AUDIO_THRESHOLD_SECONDS * sampling_rate) |
| silence_samples = int(_SHORT_AUDIO_SILENCE_SECONDS * sampling_rate) |
|
|
| return [ |
| np.pad(audio, (silence_samples, 0), mode="constant") |
| if len(audio) < threshold_samples |
| else audio |
| for audio in audios |
| ] |
|
|
|
|
| class AprielVoiceProcessor(transformers.ProcessorMixin): |
| """Process audio and images for the AprielVoice multimodal model. |
| |
| Whisper features use Apriel's audio projector. Images use the wrapped |
| Gemma4 vision tower, with one placeholder per actual image soft token. |
| """ |
|
|
| attributes = ["audio_processor", "tokenizer", "image_processor"] |
| audio_processor_class = ("WhisperFeatureExtractor",) |
| tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast") |
| image_processor_class = ("Gemma4ImageProcessor", "Gemma4ImageProcessorPil") |
|
|
| tokenizer: transformers.PreTrainedTokenizerBase |
| audio_processor: transformers.FeatureExtractionMixin |
|
|
| def __init__( |
| self, |
| audio_processor=None, |
| tokenizer=None, |
| audio_padding: str = "longest", |
| encoder_ds_factor: int = 2, |
| stack_factor: int = 8, |
| audio_placeholder: str = "<|audio|>", |
| audio_context_size: Optional[int] = 3000, |
| image_processor=None, |
| ): |
| self.audio_padding = audio_padding |
| self.encoder_ds_factor = encoder_ds_factor |
| self.stack_factor = stack_factor |
| self.audio_placeholder = audio_placeholder |
| self.audio_context_size = audio_context_size |
| assert tokenizer.eos_token is not None, "Tokenizer has no EOS token." |
| self.vocab = tokenizer.get_vocab() |
| self.audio_token_replacement = tokenizer.eos_token |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token_id = tokenizer.eos_token_id |
| if audio_processor is None: |
| audio_processor = transformers.WhisperFeatureExtractor() |
| if image_processor is None: |
| image_processor = transformers.Gemma4ImageProcessor() |
| self.image_token = getattr(tokenizer, "image_token", "<|image|>") |
| self.image_token_id = self.vocab.get(self.image_token) |
| self.boi_token = getattr(tokenizer, "boi_token", "<|image>") |
| self.eoi_token = getattr(tokenizer, "eoi_token", "<image|>") |
| super().__init__( |
| audio_processor=audio_processor, |
| tokenizer=tokenizer, |
| image_processor=image_processor, |
| ) |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs): |
| image_kwargs = dict(kwargs.pop("images_kwargs", {}) or {}) |
| for key in transformers.Gemma4ImageProcessor.valid_kwargs.__annotations__: |
| if key in kwargs: |
| image_kwargs[key] = kwargs.pop(key) |
| config: AprielVoiceConfig = transformers.AutoConfig.from_pretrained( |
| pretrained_model_name_or_path, **kwargs |
| ) |
| audio_processor = transformers.WhisperFeatureExtractor.from_pretrained( |
| pretrained_model_name_or_path, **kwargs |
| ) |
| tokenizer = transformers.AutoTokenizer.from_pretrained( |
| pretrained_model_name_or_path, **kwargs |
| ) |
| tokenizer.padding_side = "left" |
| tokenizer.pad_token = tokenizer.eos_token |
| vision_config = getattr(config.text_config, "vision_config", None) |
| image_defaults = {} |
| if vision_config is not None: |
| image_defaults = { |
| "patch_size": vision_config.patch_size, |
| "pooling_kernel_size": vision_config.pooling_kernel_size, |
| "max_soft_tokens": getattr(vision_config, "default_output_length", 280), |
| } |
| image_processor = transformers.Gemma4ImageProcessor( |
| **{**image_defaults, **image_kwargs} |
| ) |
| return cls( |
| audio_processor=audio_processor, |
| tokenizer=tokenizer, |
| stack_factor=config.stack_factor, |
| image_processor=image_processor, |
| ) |
|
|
| def _auto_chunk_audios( |
| self, |
| text: str, |
| audios: List[np.ndarray], |
| sampling_rate: int, |
| max_chunk_duration: float = 30.0, |
| ) -> Tuple[str, List[np.ndarray]]: |
| max_samples = int(max_chunk_duration * sampling_rate) |
| new_audios: List[np.ndarray] = [] |
| chunk_counts: List[int] = [] |
|
|
| for audio in audios: |
| if len(audio) > max_samples: |
| n_chunks = math.ceil(len(audio) / max_samples) |
| for i in range(n_chunks): |
| new_audios.append(audio[i * max_samples : (i + 1) * max_samples]) |
| chunk_counts.append(n_chunks) |
| else: |
| new_audios.append(audio) |
| chunk_counts.append(1) |
|
|
| if all(c == 1 for c in chunk_counts): |
| return text, audios |
|
|
| audio_block_re = re.compile( |
| r"\[AUDIO FILE \d+ START\]\s*<\|audio>\s*<\|audio\|>\s*<audio\|>\s*\[AUDIO FILE \d+ END\]" |
| ) |
| blocks = list(audio_block_re.finditer(text)) |
|
|
| if len(blocks) == len(chunk_counts): |
| new_text_parts: List[str] = [] |
| prev_end = 0 |
| file_idx = 1 |
| for block, count in zip(blocks, chunk_counts): |
| new_text_parts.append(text[prev_end:block.start()]) |
| for j in range(count): |
| if j > 0: |
| new_text_parts.append(" ") |
| new_text_parts.append( |
| f"[AUDIO FILE {file_idx} START] <|audio> <|audio|> <audio|> [AUDIO FILE {file_idx} END]" |
| ) |
| file_idx += 1 |
| prev_end = block.end() |
| new_text_parts.append(text[prev_end:]) |
| text = "".join(new_text_parts) |
| else: |
| parts = text.split("<|audio|>") |
| new_parts: List[str] = [] |
| audio_idx = 0 |
| for i, part in enumerate(parts): |
| new_parts.append(part) |
| if i < len(parts) - 1: |
| count = chunk_counts[audio_idx] if audio_idx < len(chunk_counts) else 1 |
| new_parts.append("<|audio|>" * count) |
| audio_idx += 1 |
| text = "".join(new_parts) |
|
|
| return text, new_audios |
|
|
| def _chunk_and_pad_audio( |
| self, |
| audio_values: torch.Tensor, |
| audio_lens: torch.Tensor, |
| include_audio_num_chunks: bool = False, |
| ) -> Dict[str, Any]: |
| chunked_audio_values: List[torch.Tensor] = [] |
| chunked_audio_lens: List[int] = [] |
| is_continuation_list: List[bool] = [] |
| num_chunks: List[int] = [] |
| context_size = self.audio_context_size or audio_values.shape[-1] |
|
|
| for i in range(audio_values.shape[0]): |
| num_chunks.append((int(audio_lens[i]) + context_size - 1) // context_size) |
| for offset in range(0, audio_lens[i], context_size): |
| is_continuation = offset > 0 |
| chunk = audio_values[i, :, offset : offset + context_size] |
| if is_continuation and chunk.shape[-1] < context_size: |
| chunk = F.pad(chunk, (0, context_size - chunk.shape[-1])) |
| chunked_audio_values.append(chunk) |
| chunked_audio_lens.append(min(int(audio_lens[i].item()) - offset, context_size)) |
| is_continuation_list.append(is_continuation) |
|
|
| data = { |
| "audio_values": torch.stack(chunked_audio_values, dim=0), |
| "audio_lens": torch.tensor( |
| chunked_audio_lens, dtype=torch.int64, device=audio_values.device |
| ), |
| "audio_is_continuation": torch.tensor( |
| is_continuation_list, dtype=torch.bool, device=audio_values.device |
| ), |
| "audio_batch_size": torch.tensor( |
| [len(chunked_audio_values)], device=audio_values.device |
| ), |
| } |
| if include_audio_num_chunks: |
| data["audio_num_chunks"] = torch.tensor( |
| num_chunks, dtype=torch.int64, device=audio_values.device |
| ) |
| return data |
|
|
| def __call__( |
| self, |
| text: Optional[str] = None, |
| audio: Optional[Union[np.ndarray, torch.Tensor]] = None, |
| audios: Optional[ |
| Union[List[Union[np.ndarray, torch.Tensor]], Union[np.ndarray, torch.Tensor]] |
| ] = None, |
| sampling_rate: Optional[int] = None, |
| return_tensors: Optional[ |
| Union[str, transformers.TensorType] |
| ] = transformers.TensorType.PYTORCH, |
| include_audio_num_chunks: bool = False, |
| images=None, |
| images_kwargs: Optional[Dict[str, Any]] = None, |
| **kwargs, |
| ) -> transformers.BatchFeature: |
| if text is not None and not isinstance(text, str): |
| raise ValueError("Text must be a string. Batch mode not supported yet.") |
| if audio is not None and audios is not None: |
| raise ValueError("Only one of `audio` or `audios` should be provided.") |
| if audio is not None: |
| audios = audio |
| if audios is None: |
| audios = [] |
| elif isinstance(audios, (np.ndarray, torch.Tensor)): |
| audios = [audios] if audios.ndim == 1 else list(audios) |
|
|
| if text is not None and text.count(self.audio_placeholder) != len(audios): |
| raise ValueError( |
| f"Expected {len(audios)} audio placeholders, " |
| f"found {text.count(self.audio_placeholder)}." |
| ) |
|
|
| image_kwargs = dict(images_kwargs or {}) |
| for key in self.image_processor.valid_kwargs.__annotations__: |
| if key in kwargs: |
| image_kwargs[key] = kwargs.pop(key) |
| nested_text_kwargs = kwargs.pop("text_kwargs", {}) or {} |
| nested_audio_kwargs = kwargs.pop("audio_kwargs", {}) or {} |
| text_kwargs = {**kwargs, **nested_text_kwargs} |
| audio_kwargs = {**kwargs, **nested_audio_kwargs} |
|
|
| data = {} |
| image_token_counts = [] |
| has_images = images is not None and not ( |
| isinstance(images, (list, tuple)) and len(images) == 0 |
| ) |
| if has_images: |
| if any( |
| token not in self.vocab |
| for token in (self.image_token, self.boi_token, self.eoi_token) |
| ): |
| raise ValueError("The tokenizer does not define Gemma4 image tokens.") |
| image_kwargs["return_tensors"] = "pt" |
| image_inputs = self.image_processor(images=images, **image_kwargs) |
| image_token_counts = image_inputs.pop("num_soft_tokens_per_image").tolist() |
| data.update(image_inputs) |
|
|
| if text is not None: |
| num_placeholders = text.count(self.image_token) |
| if num_placeholders != len(image_token_counts): |
| raise ValueError( |
| f"Expected {len(image_token_counts)} image placeholders, " |
| f"found {num_placeholders}." |
| ) |
| replacements = iter( |
| self.boi_token + self.image_token * count + self.eoi_token |
| for count in image_token_counts |
| ) |
| |
| text = re.sub(re.escape(self.image_token), lambda _: next(replacements), text) |
|
|
| audio_is_continuation = [] |
| if len(audios) > 0: |
| audios = [ |
| x.detach().cpu().numpy() if isinstance(x, torch.Tensor) else x |
| for x in audios |
| ] |
| sr = sampling_rate or self.audio_processor.sampling_rate |
| audios = _prepend_silence_to_short_audios(audios, sr) |
| if text is not None and not include_audio_num_chunks: |
| text, audios = self._auto_chunk_audios(text, audios, sr) |
| hop_length = self.audio_processor.hop_length |
| audios = [ |
| np.pad(x, (0, 2 * hop_length - len(x)), mode="constant") |
| if len(x) < 2 * hop_length |
| else x |
| for x in audios |
| ] |
| audio_kwargs.pop("add_special_tokens", None) |
| audio_kwargs.update( |
| sampling_rate=sr, |
| padding=self.audio_padding, |
| pad_to_multiple_of=hop_length, |
| truncation=False, |
| return_attention_mask=True, |
| ) |
| x: transformers.BatchFeature = self.audio_processor(audios, **audio_kwargs) |
| data.update( |
| self._chunk_and_pad_audio( |
| audio_values=torch.as_tensor( |
| x.input_features if "input_features" in x else x.input_values |
| ), |
| audio_lens=torch.as_tensor(x.attention_mask).sum(-1), |
| include_audio_num_chunks=include_audio_num_chunks, |
| ) |
| ) |
| audio_is_continuation = data.pop("audio_is_continuation") |
| data["audio_token_len"] = torch.ceil( |
| data["audio_lens"] / (self.encoder_ds_factor * self.stack_factor) |
| ).to(dtype=torch.int) |
|
|
| if text is not None: |
| text_kwargs["add_special_tokens"] = False |
| tokenized_parts = self.tokenizer( |
| text.split(self.audio_placeholder), |
| **text_kwargs, |
| ) |
|
|
| audio_token_start_idx = [] |
| placeholder_index = -1 |
| split_input_ids = tokenized_parts["input_ids"] |
| input_ids: List[int] = [] |
| audio_replacement_token_id = self.vocab[self.audio_token_replacement] |
|
|
| for i, token_len in enumerate(data.get("audio_token_len", [])): |
| if not audio_is_continuation[i]: |
| placeholder_index += 1 |
| if placeholder_index >= len(split_input_ids): |
| raise ValueError(f"Too few audio placeholders (expected {len(audios)}).") |
| input_ids.extend(split_input_ids[placeholder_index]) |
| audio_token_start_idx.append(len(input_ids)) |
| input_ids.extend([audio_replacement_token_id] * token_len) |
|
|
| placeholder_index += 1 |
| if placeholder_index != len(split_input_ids) - 1: |
| raise ValueError(f"Too many audio placeholders (expected {len(audios)}).") |
| input_ids.extend(split_input_ids[placeholder_index]) |
|
|
| if "audio_token_len" in data: |
| data["audio_token_start_idx"] = torch.as_tensor(audio_token_start_idx) |
|
|
| data["input_ids"] = [input_ids] |
| data["attention_mask"] = [[1] * len(input_ids)] |
| if image_token_counts: |
| if input_ids.count(self.image_token_id) != sum(image_token_counts): |
| raise ValueError("Tokenization changed the number of image placeholders.") |
| mm_token_type_ids = [ |
| int(token == self.image_token_id) for token in input_ids |
| ] |
| for start, length in zip( |
| audio_token_start_idx, data.get("audio_token_len", []) |
| ): |
| mm_token_type_ids[start : start + int(length)] = [3] * int(length) |
| data["mm_token_type_ids"] = [mm_token_type_ids] |
|
|
| return transformers.BatchFeature(data=data, tensor_type=return_tensors) |
|
|
| def batch_decode(self, *args, **kwargs): |
| return self.tokenizer.batch_decode(*args, **kwargs) |
|
|
| def decode(self, *args, **kwargs): |
| return self.tokenizer.decode(*args, **kwargs) |
|
|
| @property |
| def model_input_names(self): |
| return list(dict.fromkeys(self.tokenizer.model_input_names + [ |
| "audio_values", "audio_lens", "audio_token_len", "audio_token_start_idx", |
| "audio_batch_size", "pixel_values", "image_position_ids", "mm_token_type_ids", |
| ])) |
|
|
|
|
| AprielVoiceProcessor.register_for_auto_class() |
| transformers.AutoProcessor.register(AprielVoiceConfig, AprielVoiceProcessor) |
|
|
