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

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=1, 
        chunk_length=30, 
        overlap_seconds=10, 
        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.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)
        # 注意：这里我们只处理不带重叠的块，重叠将在外部处理（如果需要）
        # 或者在迭代器内部更明确地处理。为了简化，我们假设分块是基于 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 = torch.zeros(self.batch_size, dtype=torch.long)
        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)
            if padding_size > 0:
                output_lengths[-1] = kernel - padding_size
            output_seq_no = torch.full((num_chunks,), seq_no, dtype=torch.long)
            return output_tensor, output_lengths, output_seq_no

        for i, sample in enumerate(self.data):
            sample_chunks, sample_lengths, sample_seq_no = 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[start_idx_batch:end_idx_batch] = 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 = [
                        wav_tensor[xi, :, :x_len].reshape(-1).cpu().numpy() 
                        for xi, x_len in enumerate(input_lengths.tolist())
                    ]
                    yield BatchFeature({
                        **self.encode_func(list_x),
                        "chunk_seq_no": input_seq_no.clone(),    
                    })
                    
                    # 重置批次计数器和Tensor内容
                    batch_num = 0
                    wav_tensor.zero_()
                    input_lengths.zero_()
                    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())
            ]
            yield BatchFeature({
                **self.encode_func(list_x),
                "chunk_seq_no": input_seq_no[:batch_num].clone(),    
            })


class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
    def __init__(
        self,
        feature_size=80,
        sampling_rate=16000,
        hop_length=160,
        chunk_length=30,
        n_fft=400,
        padding_value=0.0,
        dither=0.0,
        return_attention_mask=False,
        max_frequency=None,
        batch_size=None,
        **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,
            **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",
        )

    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=len(raw_speech) if self.batch_size is None else self.batch_size, 
            chunk_length=self.chunk_length, 
            overlap_seconds=overlap_seconds, 
            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,
            )
        )