feature_extraction_xy_tokenizer.py

# coding=utf-8
# Copyright 2022 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
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
from transformers import WhisperFeatureExtractor
from transformers.audio_utils import mel_filter_bank
from transformers.configuration_utils import PretrainedConfig
from transformers.feature_extraction_utils import BatchFeature
from transformers.utils import TensorType, logging

logger = logging.get_logger(__name__)


class ExtractorIterator:
    def __init__(
        self, 
        data, 
        batch_size=8, 
        chunk_length=30, 
        overlap_seconds=10,
        overlap_side="both", 
        sampling_rate=16000,
        encode_func = None,
    ) -> None:
        self.data = data
        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.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
        
        assert callable(encode_func)
        self.encode_func = encode_func

    def __iter__(self):
        """
        返回一个生成器，该生成器负责处理所有批处理逻辑。
        这是最 Pythonic 的实现方式。
        """
        # 批处理相关的变量现在是 __iter__ 的局部变量，非常清晰
        batch_num = 0
        
        # 注意：chunk_and_pad_view 输出的块大小是 duration_size
        wav_tensor = torch.zeros(self.batch_size, 1, self.chunk_size)
        input_lengths = deque(maxlen=self.batch_size)
        input_seq_no = torch.zeros(self.batch_size, dtype=torch.long)
        
        right_boundary = self.get_right_boundary()

        for i, sample in enumerate(self.data):
            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):
                space_in_batch = self.batch_size - batch_num
                chunks_to_add = min(space_in_batch, len(sample_chunks) - processed_in_sample)
                
                # 定义切片范围
                start_idx_sample = processed_in_sample
                end_idx_sample = processed_in_sample + chunks_to_add
                start_idx_batch = batch_num
                end_idx_batch = batch_num + chunks_to_add
                
                # 填充数据
                wav_tensor[start_idx_batch:end_idx_batch] = sample_chunks[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]

                # 更新计数器
                batch_num += chunks_to_add
                processed_in_sample += chunks_to_add

                # 如果批次满了，yield 一个副本并重置
                if batch_num == self.batch_size:
                    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,
                        "chunk_seq_no": input_seq_no.clone(),    
                    })
                    
                    # 重置批次计数器和Tensor内容
                    batch_num = 0
                    wav_tensor.zero_()
                    input_lengths.clear()
                    input_seq_no.zero_()

        # 循环结束后，处理最后一个未满的批次
        if batch_num > 0:
            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__(
        self,
        feature_size=80,
        sampling_rate=16000,
        hop_length=160,
        chunk_length=30,
        n_fft=400,
        n_samples=480000,
        nb_max_frames=3000,
        padding_side="right",
        padding_value=0.0,
        dither=0.0,
        return_attention_mask=False,
        max_frequency=None,
        batch_size=8,
        overlap_side="both", 
        **kwargs,
    ):
        super().__init__(
            feature_size=feature_size,
            sampling_rate=sampling_rate,
            hop_length=hop_length,
            chunk_length=chunk_length,
            n_fft=n_fft,
            padding_value=padding_value,
            dither=dither,
            return_attention_mask=return_attention_mask,
            n_samples=n_samples,
            nb_max_frames=nb_max_frames,
            padding_side=padding_side,
            **kwargs,
        )
        self.max_frequency = max_frequency if max_frequency is not None else sampling_rate / 2
        self.batch_size = batch_size
        self.mel_filters = mel_filter_bank(
            num_frequency_bins=1 + n_fft // 2,
            num_mel_filters=feature_size,
            min_frequency=0.0,
            max_frequency=self.max_frequency,
            sampling_rate=sampling_rate,
            norm="slaney",
            mel_scale="slaney",
        )
        self.overlap_side = overlap_side

    def __call__(
        self,
        raw_speech: Union[torch.Tensor, List[torch.Tensor]],
        truncation: bool = True,
        pad_to_multiple_of: Optional[int] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        return_attention_mask: Optional[bool] = None,
        padding: Optional[str] = "max_length",
        max_length: Optional[int] = None,
        sampling_rate: Optional[int] = None,
        do_normalize: Optional[bool] = None,
        device: Optional[str] = "cpu",
        return_token_timestamps: Optional[bool] = None,
        overlap_seconds: int = 10,
        **kwargs,
    ) -> ExtractorIterator:

        if not isinstance(raw_speech, list):
            raw_speech = [raw_speech]
    
        return ExtractorIterator(
            raw_speech,
            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__,
                truncation=truncation,
                pad_to_multiple_of=pad_to_multiple_of,
                return_tensors=return_tensors,
                return_attention_mask=return_attention_mask,
                padding=padding,
                max_length=max_length,
                sampling_rate=sampling_rate,
                do_normalize=do_normalize,
                device=device,
                return_token_timestamps=return_token_timestamps,
                **kwargs,
            )
        )