main
py 229 lines 8.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # ruff: noqa
5 # mypy: ignore-errors
6 # coding=utf-8
7 # Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
8 #
9 # Licensed under the Apache License, Version 2.0 (the "License");
10 # you may not use this file except in compliance with the License.
11 # You may obtain a copy of the License at
12 #
13 # http://www.apache.org/licenses/LICENSE-2.0
14 #
15 # Unless required by applicable law or agreed to in writing, software
16 # distributed under the License is distributed on an "AS IS" BASIS,
17 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18 # See the License for the specific language governing permissions and
19 # limitations under the License.
20 import regex as re
21
22 import numpy as np
23
24 from transformers import AutoProcessor
25 from transformers.audio_utils import AudioInput
26 from transformers.feature_extraction_utils import BatchFeature
27 from transformers.processing_utils import ProcessingKwargs, ProcessorMixin
28 from transformers.tokenization_utils_base import TextInput
29
30
31 class Qwen3ASRProcessorKwargs(ProcessingKwargs, total=False):
32 _defaults = {
33 "text_kwargs": {
34 "padding": False,
35 "padding_side": "left",
36 },
37 "audio_kwargs": {
38 "sampling_rate": 16000,
39 "padding": True,
40 "return_attention_mask": True,
41 },
42 }
43
44
45 def _get_feat_extract_output_lengths(input_lengths):
46 """
47 Computes the output length of the convolutional layers and the output length of the audio encoder
48 """
49
50 input_lengths_leave = input_lengths % 100
51 feat_lengths = (input_lengths_leave - 1) // 2 + 1
52 output_lengths = (
53 ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
54 )
55 return output_lengths
56
57
58 class Qwen3ASRProcessor(ProcessorMixin):
59 r"""
60 Constructs a Qwen3ASR processor.
61 [`Qwen3ASRProcessor`] offers all the functionalities of [`WhisperFeatureExtractor`], and [`Qwen2TokenizerFast`]. See the
62 [`~Qwen3ASRProcessor.__call__`] and [`~Qwen3ASRProcessor.decode`] for more information.
63
64 Args:
65 feature_extractor ([`WhisperFeatureExtractor`], *optional*):
66 The audio feature extractor.
67 tokenizer ([`Qwen2TokenizerFast`], *optional*):
68 The text tokenizer.
69 chat_template (`Optional[str]`, *optional*):
70 The Jinja template to use for formatting the conversation. If not provided, the default chat template is used.
71 """
72
73 attributes = ["feature_extractor", "tokenizer"]
74 feature_extractor_class = "WhisperFeatureExtractor"
75 tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
76
77 def __init__(self, feature_extractor=None, tokenizer=None, chat_template=None):
78 super().__init__(feature_extractor, tokenizer, chat_template=chat_template)
79 self.audio_token = self.tokenizer.audio_token
80 self.audio_bos_token = self.tokenizer.audio_bos_token
81 self.audio_eos_token = self.tokenizer.audio_eos_token
82
83 def __call__(
84 self,
85 text: TextInput = None,
86 audio: AudioInput = None,
87 **kwargs,
88 ) -> BatchFeature:
89 """
90 Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text`
91 and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
92 the text. To prepare the audio(s), this method forwards the `audio` and `kwargs` arguments to
93 WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audio` is not `None`. Please refer to the doctsring
94 of the above two methods for more information.
95
96 Args:
97 text (`str`, `List[str]`, `List[List[str]]`):
98 The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
99 (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
100 `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
101 audio (`np.ndarray`, `List[np.ndarray]`):
102 The audio or batch of audio to be prepared. Each audio can be a NumPy array.
103 """
104
105 if text is None:
106 raise ValueError("You need to specify either a `text` input to process.")
107
108 output_kwargs = self._merge_kwargs(
109 Qwen3ASRProcessorKwargs,
110 tokenizer_init_kwargs=self.tokenizer.init_kwargs,
111 **kwargs,
112 )
113
114 if audio is not None:
115 output_kwargs["audio_kwargs"]["padding"] = True
116 output_kwargs["audio_kwargs"]["truncation"] = False
117 audio_inputs = self.feature_extractor(
118 audio, **output_kwargs["audio_kwargs"]
119 )
120 audio_inputs["feature_attention_mask"] = audio_inputs.pop(
121 "attention_mask"
122 ) # rename feature_attention_mask to prevent conflicts later on
123 audio_inputs["input_features"] = audio_inputs.pop(
124 "input_features"
125 ) # rename input_features to prevent conflicts later on
126 audio_lengths = iter(
127 _get_feat_extract_output_lengths(
128 audio_inputs["feature_attention_mask"].sum(-1)
129 )
130 )
131 else:
132 audio_inputs = {}
133 audio_lengths = iter([])
134
135 if not isinstance(text, list):
136 text = [text]
137
138 text = self.replace_multimodal_special_tokens(
139 text,
140 audio_lengths,
141 )
142
143 texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
144
145 return BatchFeature(
146 data={**texts_inputs, **audio_inputs},
147 tensor_type=kwargs.get("return_tensors"),
148 )
149
150 def replace_multimodal_special_tokens(
151 self,
152 text,
153 audio_lengths,
154 ):
155 processed_text = []
156 for sample in text:
157 positions = []
158 special_tokens = [re.escape(tok) for tok in [self.audio_token]]
159 pattern = "|".join(special_tokens)
160 positions = sorted(
161 [
162 (match.start(), match.group())
163 for match in re.finditer(pattern, sample)
164 ]
165 )
166 positions.sort(key=lambda x: x[0])
167
168 for _, special_token in positions:
169 if special_token == self.audio_token:
170 sample = sample.replace(
171 self.audio_token,
172 "<|audio_placeholder|>" * next(audio_lengths),
173 1,
174 )
175
176 sample = sample.replace("<|audio_placeholder|>", self.audio_token)
177 processed_text.append(sample)
178 return processed_text
179
180 def get_chunked_index(
181 self, token_indices: np.ndarray, tokens_per_chunk: int
182 ) -> list[tuple[int, int]]:
183 """
184 Splits token index list into chunks based on token value ranges.
185
186 Given a list of token indices, returns a list of (start, end) index tuples representing
187 slices of the list where the token values fall within successive ranges of `tokens_per_chunk`.
188
189 For example, if `tokens_per_chunk` is 1000, the function will create chunks such that:
190 - the first chunk contains token values < 1000,
191 - the second chunk contains values >= 1000 and < 2000, and so on.
192
193 Parameters:
194 token_indices (`np.ndarray`): A monotonically increasing list of token index values.
195 tokens_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
196
197 Returns:
198 `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
199 and end (exclusive) indices of a chunk in `token_indices`.
200 """
201
202 def _iter():
203 i, start_idx = 0, 0 # skip bos token
204 current_chunk = 1
205 while i < len(token_indices): # skip eos token
206 if token_indices[i] >= current_chunk * tokens_per_chunk:
207 yield (start_idx, i)
208 start_idx = i
209 current_chunk += 1
210 i += 1
211 yield (start_idx, len(token_indices))
212
213 return list(_iter())
214
215 def apply_chat_template(self, conversations, chat_template=None, **kwargs):
216 kwargs["return_dict"] = False
217 return super().apply_chat_template(conversations, chat_template, **kwargs)
218
219 @property
220 def model_input_names(self):
221 tokenizer_input_names = self.tokenizer.model_input_names
222 feature_extractor_input_names = self.feature_extractor.model_input_names
223 return list(
224 dict.fromkeys(
225 tokenizer_input_names
226 + feature_extractor_input_names
227 + ["feature_attention_mask"]
228 )
229 )