Upload processing_borealis.py with huggingface_hub

processing_borealis.py

@@ -4,6 +4,7 @@ Borealis Processor for HuggingFace/vLLM compatibility.
 Handles audio feature extraction and tokenization.
 """
 
+import numpy as np
 from typing import List, Optional, Union
 
 import torch
@@ -28,6 +29,9 @@ class BorealisProcessor(ProcessorMixin):
     audio_bos_token = "<|start_of_audio|>"
     audio_eos_token = "<|start_of_audio|>"  # Reuse bos token since only 2 audio tokens in vocab
 
+    # Borealis architecture parameters
+    downsample_factor = 4  # Audio embedding downsampling factor
+
     def __init__(
         self,
         feature_extractor: Optional[WhisperFeatureExtractor] = None,
@@ -59,8 +63,10 @@ class BorealisProcessor(ProcessorMixin):
         """
         Process text and/or audio inputs.
 
+        Expands <|AUDIO|> tokens in text to match the number of audio embeddings.
+
         Args:
-            text: Text prompt(s)
+            text: Text prompt(s) containing <|AUDIO|> placeholders
             audio: Audio waveform(s) at 16kHz
             audios: Audio waveform(s) at 16kHz (vLLM style)
             sampling_rate: Audio sampling rate (default: 16000)
@@ -88,20 +94,62 @@ class BorealisProcessor(ProcessorMixin):
             for a in audio:
                 if isinstance(a, torch.Tensor):
                     a = a.numpy()
+                if isinstance(a, np.ndarray):
+                    a = a.astype(np.float32)
                 audio_arrays.append(a)
 
             audio_features = self.feature_extractor(
                 audio_arrays,
                 sampling_rate=sampling_rate,
                 return_tensors=return_tensors,
+                padding="max_length",
+                return_attention_mask=True,
             )
             data["input_features"] = audio_features.input_features
 
-        # Process text if provided
+            # Calculate audio lengths for token expansion
+            # Whisper uses 30s chunks with 3000 mel frames -> 1500 encoder frames
+            # Borealis downsamples by 4x -> 375 tokens
+            attention_mask = audio_features.get("attention_mask")
+            if attention_mask is not None:
+                # Sum attention mask to get actual audio length in frames
+                audio_lengths = attention_mask.sum(dim=-1).tolist()
+            else:
+                # Default: assume full 30s audio
+                audio_lengths = [3000] * len(audio_arrays)
+
+        # Process text if provided - expand audio tokens
         if text is not None:
             if isinstance(text, str):
                 text = [text]
 
+            # Expand <|AUDIO|> tokens based on audio lengths
+            if audio is not None:
+                expanded_text = []
+                audio_idx = 0
+
+                for sample in text:
+                    while self.audio_token in sample:
+                        if audio_idx < len(audio_lengths):
+                            audio_len = audio_lengths[audio_idx]
+                            # Whisper: 3000 mel frames -> 1500 encoder frames
+                            # Then downsample by 4 -> 375 tokens
+                            whisper_frames = (audio_len - 1) // 2 + 1  # ~1500
+                            num_audio_tokens = whisper_frames // self.downsample_factor  # ~375
+
+                            # Expand single <|AUDIO|> to multiple tokens with markers
+                            expanded = (
+                                self.audio_bos_token +
+                                self.audio_token * num_audio_tokens +
+                                self.audio_eos_token
+                            )
+                            sample = sample.replace(self.audio_token, expanded, 1)
+                            audio_idx += 1
+                        else:
+                            break
+                    expanded_text.append(sample)
+                text = expanded_text
+
             # Filter out kwargs that tokenizer doesn't accept
             tok_kwargs = {k: v for k, v in kwargs.items()
                          if k in ['padding', 'truncation', 'max_length', 'add_special_tokens']}