add middle overlap type

feature_extraction_xy_tokenizer.py

@@ -18,6 +18,7 @@ Feature extractor class for Whisper
 import math
 from functools import partial
 from typing import List, Optional, Union
+from collections import deque
 
 import torch
 import torch.nn.functional as F
@@ -34,9 +35,10 @@ class ExtractorIterator:
     def __init__(
         self, 
         data, 
-        batch_size=1, 
+        batch_size=8, 
         chunk_length=30, 
-        overlap_seconds=10, 
+        overlap_seconds=10,
+        overlap_side="both", 
         sampling_rate=16000,
         encode_func = None,
     ) -> None:
@@ -44,12 +46,16 @@ class ExtractorIterator:
         self.batch_size = batch_size
         self.chunk_length = chunk_length
         self.overlap_seconds = overlap_seconds
+        self.overlap_side = overlap_side
         self.sampling_rate = sampling_rate
         
         # duration_size 是每次处理的有效音频长度
         self.chunk_size = int(self.chunk_length * self.sampling_rate)
-        self.duration_seconds = self.chunk_length - self.overlap_seconds
-        self.duration_size = int(self.duration_seconds * self.sampling_rate)
+        self.overlap_size = int(self.overlap_seconds * self.sampling_rate)
+        self.duration_size = self.chunk_size - self.overlap_size
+        assert (
+            (overlap_side == "right") or (self.overlap_size % 2 == 0)
+        ), '`overlap_seconds` must be divisible by 2 when `overlap_side` is "both".'
         # 注意：这里我们只处理不带重叠的块，重叠将在外部处理（如果需要）
         # 或者在迭代器内部更明确地处理。为了简化，我们假设分块是基于 duration_size
         
@@ -66,29 +72,13 @@ class ExtractorIterator:
         
         # 注意：chunk_and_pad_view 输出的块大小是 duration_size
         wav_tensor = torch.zeros(self.batch_size, 1, self.chunk_size)
-        input_lengths = torch.zeros(self.batch_size, dtype=torch.long)
+        input_lengths = deque(maxlen=self.batch_size)
         input_seq_no = torch.zeros(self.batch_size, dtype=torch.long)
-
-        def chunk_and_pad_view(tensor, seq_no):
-            x = tensor[0:1, :].unsqueeze(0)
-            
-            stride = self.duration_size
-            kernel = self.chunk_size
-            B, C, L = x.shape
-            
-            num_chunks = math.ceil(L / stride)
-            target_len = (num_chunks - 1) * stride + kernel
-            padding_size = max(0, target_len - L)
-            x_padded = F.pad(x, (0, padding_size), "constant", 0)
-            output_tensor = x_padded.unfold(dimension=2, size=kernel, step=stride).squeeze(0).transpose(0, 1)
-            output_lengths = torch.full((num_chunks,), kernel, dtype=torch.long)
-            for i in range(num_chunks):
-                output_lengths[i] = min(output_lengths[i], L - stride * i)
-            output_seq_no = torch.full((num_chunks,), seq_no, dtype=torch.long)
-            return output_tensor, output_lengths, output_seq_no
+        
+        right_boundary = self.get_right_boundary()
 
         for i, sample in enumerate(self.data):
-            sample_chunks, sample_lengths, sample_seq_no = chunk_and_pad_view(sample, i)
+            sample_chunks, sample_lengths, sample_seq_no = self.chunk_and_pad_view(sample, i)
             
             processed_in_sample = 0
             while processed_in_sample < len(sample_chunks):
@@ -103,7 +93,7 @@ class ExtractorIterator:
                 
                 # 填充数据
                 wav_tensor[start_idx_batch:end_idx_batch] = sample_chunks[start_idx_sample:end_idx_sample]
-                input_lengths[start_idx_batch:end_idx_batch] = sample_lengths[start_idx_sample:end_idx_sample]
+                input_lengths.extend(sample_lengths[start_idx_sample:end_idx_sample])
                 input_seq_no[start_idx_batch:end_idx_batch] = sample_seq_no[start_idx_sample:end_idx_sample]
 
                 # 更新计数器
@@ -112,10 +102,13 @@ class ExtractorIterator:
 
                 # 如果批次满了，yield 一个副本并重置
                 if batch_num == self.batch_size:
-                    list_x = [
-                        wav_tensor[xi, :, :x_len].reshape(-1).cpu().numpy() 
-                        for xi, x_len in enumerate(input_lengths.tolist())
-                    ]
+                    list_x = []
+                    for xi, (_, right) in enumerate(input_lengths):
+                        if right == right_boundary and torch.any(wav_tensor[xi, :, right:] != 0):
+                            list_x.append(wav_tensor[xi].reshape(-1).cpu().numpy())
+                        else:
+                            list_x.append(wav_tensor[xi, :, :right].reshape(-1).cpu().numpy())
+                            
                     yield BatchFeature({
                         **self.encode_func(list_x),
                         "input_lengths": input_lengths,
@@ -125,21 +118,62 @@ class ExtractorIterator:
                     # 重置批次计数器和Tensor内容
                     batch_num = 0
                     wav_tensor.zero_()
-                    input_lengths.zero_()
+                    input_lengths.clear()
                     input_seq_no.zero_()
 
         # 循环结束后，处理最后一个未满的批次
         if batch_num > 0:
-            list_x = [
-                wav_tensor[xi, :, :x_len].reshape(-1).cpu().numpy() 
-                for xi, x_len in enumerate(input_lengths.tolist())
-            ]
+            list_x = []
+            for xi in range(batch_num):
+                _, right = input_lengths[xi]
+                if right == right_boundary and torch.any(wav_tensor[xi, :, right:] != 0):
+                    list_x.append(wav_tensor[xi].reshape(-1).cpu().numpy())
+                else:
+                    list_x.append(wav_tensor[xi, :, :right].reshape(-1).cpu().numpy())
             yield BatchFeature({
                 **self.encode_func(list_x),
                 "input_lengths": input_lengths,
                 "chunk_seq_no": input_seq_no[:batch_num].clone(),    
             })
 
+    def chunk_and_pad_view(self, tensor, seq_no):
+        x = tensor[0:1, :].unsqueeze(0)
+        
+        stride = self.duration_size
+        kernel = self.chunk_size
+        B, C, L = x.shape
+    
+        num_chunks = max(0, math.ceil((L - kernel) / stride)) + 1
+        target_len = (num_chunks - 1) * stride + kernel
+        padding_size = max(0, target_len - L)
+        x_padded = F.pad(x, (0, padding_size), "constant", 0)
+        output_tensor = x_padded.unfold(dimension=2, size=kernel, step=stride).squeeze(0).transpose(0, 1)
+        
+        output_lengths = self.get_windows_boundaries(num_chunks, L)
+        output_seq_no = torch.full((num_chunks,), seq_no, dtype=torch.long)
+        return output_tensor, output_lengths, output_seq_no
+
+    def get_left_boundary(self):
+        if self.overlap_side == "right":
+            return 0
+        else:
+            return int(self.overlap_size / 2)
+
+    def get_right_boundary(self):
+        if self.overlap_side == "right":
+            return self.duration_size
+        else:
+            return self.chunk_size - int(self.overlap_size / 2)
+            
+    def get_windows_boundaries(self, num_chunks, seq_len):
+        left_boundary = self.get_left_boundary()
+        right_boundary = self.get_right_boundary()
+
+        output_lengths = [(left_boundary, right_boundary) for _ in range(num_chunks)]
+        output_lengths[0] = (0, output_lengths[0][1])
+        output_lengths[-1] = (output_lengths[-1][0], seq_len - self.duration_size * (num_chunks-1))
+        return output_lengths
+
 
 class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
     def __init__(
@@ -156,7 +190,8 @@ class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
         dither=0.0,
         return_attention_mask=False,
         max_frequency=None,
-        batch_size=None,
+        batch_size=8,
+        overlap_side="both", 
         **kwargs,
     ):
         super().__init__(
@@ -184,6 +219,7 @@ class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
             norm="slaney",
             mel_scale="slaney",
         )
+        self.overlap_side = overlap_side
 
     def __call__(
         self,
@@ -207,9 +243,10 @@ class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
     
         return ExtractorIterator(
             raw_speech,
-            batch_size=len(raw_speech) if self.batch_size is None else self.batch_size, 
+            batch_size=self.batch_size if self.batch_size else len(raw_speech), 
             chunk_length=self.chunk_length, 
             overlap_seconds=overlap_seconds, 
+            overlap_side=self.overlap_side,
             sampling_rate=self.sampling_rate,
             encode_func=partial(
                 super().__call__,

modeling_xy_tokenizer.py

@@ -858,6 +858,16 @@ class XYTokenizerModel(XYTokenizerPreTrainedModel):
             
         return torch.tensor([_get_out_len(l) for l in input_lengths], device=self.device)
 
+    def scale_window_size(self, boundaries, scaling_factor):
+        scaling_range = []
+        scaling_boundaries = []
+        for left_boundary, right_boundary in boundaries:
+            scaling_left_boundary = left_boundary// scaling_factor
+            scaling_right_boundary = right_boundary // scaling_factor
+            scaling_range.append(scaling_right_boundary-scaling_left_boundary)
+            scaling_boundaries.append(slice(scaling_left_boundary, scaling_right_boundary))
+        return scaling_range, scaling_boundaries
+
     @torch.inference_mode
     def encode(
         self,
@@ -896,11 +906,11 @@ class XYTokenizerModel(XYTokenizerPreTrainedModel):
             for chunk_features in features:
                 # Always use return_dict=True for easier access to named outputs
                 chunk_output = self._encode(chunk_features, n_quantizers, return_dict=True)
-                valid_code_lengths = torch.clamp(chunk_features["input_lengths"], 0, features.duration_size) // self.encoder_downsample_rate
+                valid_code_lengths, valid_code_ranges = self.scale_window_size(chunk_features["input_lengths"], self.encoder_downsample_rate)
 
                 # Accumulate weighted commit loss
                 chunk_length = chunk_output.codes_lengths.sum().item()
-                valid_chunk_length = valid_code_lengths.sum().item()
+                valid_chunk_length = sum(valid_code_lengths)
                 if chunk_output.commit_loss is not None and valid_chunk_length > 0:
                     commit_loss = chunk_output.commit_loss / chunk_length * valid_chunk_length
                     commit_losses.append((commit_loss.cpu(), valid_chunk_length))
@@ -908,12 +918,12 @@ class XYTokenizerModel(XYTokenizerPreTrainedModel):
 
                 # Group results by original sequence ID
                 for i, seq_id in enumerate(chunk_features["chunk_seq_no"].tolist()):
-                    valid_code_length = valid_code_lengths[i]
-                    if valid_code_length > 0:
-                        encodings[seq_id]["zq"].append(chunk_output.quantized_representation[i:i+1, :, :valid_code_length])
-                        encodings[seq_id]["codes"].append(chunk_output.audio_codes[:, i:i+1, :valid_code_length])
+                    valid_code_range = valid_code_ranges[i]
+                    if valid_code_range.stop > 0:
+                        encodings[seq_id]["zq"].append(chunk_output.quantized_representation[i:i+1, :, valid_code_range])
+                        encodings[seq_id]["codes"].append(chunk_output.audio_codes[:, i:i+1, valid_code_range])
                         # Add the valid length of this chunk to the total for this sequence
-                        encodings[seq_id]["length"] += valid_code_lengths[i].item()
+                        encodings[seq_id]["length"] += valid_code_lengths[i]
 
             final_outputs = []
             for seq_id, seq_data in encodings.items():

preprocessor_config.json

@@ -9,5 +9,6 @@
   "padding_value": 0.0,
   "sampling_rate": 16000,
   "return_attention_mask": true,
-  "return_tensors": "pt"
+  "return_tensors": "pt",
+  "overlap_side": "both"
 }