Ming-3.0-asr-tiny / processing_ming_v3_asr.py

Duplicate from inclusionAI/Ming-3.0-asr-tiny

bef6507 3 days ago

10.7 kB

# Copyright 2026 Ant Group. All rights reserved.
# 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 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.
"""Hugging Face processor for Ming V3 ASR transcription."""

import re

import numpy as np
import torch
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessingKwargs, ProcessorMixin

from .audio_spec import MingAudioSpec  # noqa: F401 - include transitive dependency in HF's local module cache
from .feature_extraction_ming_v3_asr import MingV3ASRFeatureExtractor

_AUDIO_FORMS = (
    ("<audio>", "<audioPatch>", "</audio>"),
    ("<|audio_start|>", "<|audio_pad|>", "<|audio_end|>"),
)
_PLACEHOLDER = "<AUDIO>"
_ASR_INSTRUCTION = (
    "Please recognize the language of this speech and transcribe it, "
    "placing a `\t` between the identified language and the transcribe."
)


class MingV3ASRProcessorKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "text_kwargs": {
            "padding": False,
            "padding_side": "left",
            "return_attention_mask": True,
            "return_token_type_ids": False,
        },
    }


class MingV3ASRProcessor(ProcessorMixin):
    """Expand audio tokens and combine an unchanged tokenizer with log-mel.
    ``audio`` (or the compatible ``audios`` spelling) is a flat list of clips in
    prompt-placeholder order. ``audio_num_clips`` records each text request's
    clip count, including zero. Audio features and raw lengths are flattened
    independently of the text batch. Image/video inputs are not supported.
    Chat-template support is inherited from ``ProcessorMixin`` for tokenizer
    compatibility. ASR inference uses ``apply_transcription_request`` instead.
    """

    feature_extractor_class = "AutoFeatureExtractor"
    tokenizer_class = "AutoTokenizer"
    valid_processor_kwargs = MingV3ASRProcessorKwargs

    def __init__(
        self,
        feature_extractor: MingV3ASRFeatureExtractor,
        tokenizer,
        chat_template=None,
    ):
        template = tokenizer.chat_template if chat_template is None else chat_template
        super().__init__(feature_extractor, tokenizer, chat_template=template)
        vocabulary = tokenizer.get_vocab()
        self.audio_start_token, self.audio_token, self.audio_end_token = next(
            (
                form
                for form in _AUDIO_FORMS
                if all(token in vocabulary for token in form)
            ),
            (None, None, None),
        )
        if self.audio_token is None:
            raise ValueError("tokenizer lacks Ming audio start/patch/end tokens")
        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)

    @staticmethod
    def _normalize_prompt(text):
        for start, patch, end in _AUDIO_FORMS:
            text = re.sub(
                re.escape(start) + "(?:" + re.escape(patch) + ")+" + re.escape(end),
                _PLACEHOLDER,
                text,
            )
        return text

    def __call__(
        self, text=None, audio=None, *, images=None, videos=None, audios=None, **kwargs
    ):
        if images is not None or videos is not None:
            raise ValueError("MingV3ASRProcessor supports audio and text only")
        if audio is not None and audios is not None:
            raise ValueError("Pass either audio or audios, not both")
        clips = self.feature_extractor.as_clips(audio if audio is not None else audios)
        if text is None:
            if not clips:
                raise ValueError("Provide text or audio")
            text = " ".join([_PLACEHOLDER] * len(clips))
        texts = [text] if isinstance(text, str) else text
        if (
            not isinstance(texts, (list, tuple))
            or not texts
            or not all(isinstance(t, str) for t in texts)
        ):
            raise ValueError("text must be a string or a nonempty list of strings")
        texts = [self._normalize_prompt(t) for t in texts]
        clip_counts = [t.count(_PLACEHOLDER) for t in texts]
        if sum(clip_counts) != len(clips):
            raise ValueError(
                f"Audio placeholder count {sum(clip_counts)} does not match {len(clips)} clips"
            )

        options = self._merge_kwargs(
            MingV3ASRProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )
        text_options = options["text_kwargs"]
        audio_options = options["audio_kwargs"]
        return_tensors = text_options.get("return_tensors")
        audio_return_tensors = audio_options.get("return_tensors")
        if (
            return_tensors is not None
            and audio_return_tensors is not None
            and return_tensors != audio_return_tensors
        ):
            raise ValueError("text and audio return_tensors must match")
        if return_tensors is None:
            return_tensors = audio_return_tensors
        if return_tensors is not None:
            text_options["return_tensors"] = return_tensors
            audio_options["return_tensors"] = return_tensors
        if text_options.get("return_overflowing_tokens"):
            raise ValueError("Audio token overflow windows are not supported")
        # Top-level max_length/truncation refer to text; audio is always chunked.
        # Explicit nested audio options retain extractor validation.
        for key in ("max_length", "truncation", "padding"):
            if key not in kwargs.get("audio_kwargs", {}):
                audio_options.pop(key, None)
        audio_inputs = self.feature_extractor(clips, **audio_options) if clips else {}
        lengths = audio_inputs.get("audio_feats_lengths", [])
        token_counts = [
            self.feature_extractor.spec.token_count(int(length)) for length in lengths
        ]

        expected_counts = []
        clip_index = 0
        for index, count in enumerate(clip_counts):
            expected_counts.append(sum(token_counts[clip_index : clip_index + count]))
            for token_count in token_counts[clip_index : clip_index + count]:
                replacement = (
                    self.audio_start_token
                    + self.audio_token * token_count
                    + self.audio_end_token
                )
                texts[index] = texts[index].replace(_PLACEHOLDER, replacement, 1)
            clip_index += count
        encoded = self.tokenizer(texts, **text_options)
        for ids, expected in zip(encoded["input_ids"], expected_counts):
            actual = sum(int(token) == self.audio_token_id for token in ids)
            if actual != expected:
                raise ValueError(
                    "Audio token count changed during tokenization; do not truncate audio placeholders"
                )
        if clips:
            if return_tensors == "pt":
                audio_inputs["audio_num_clips"] = torch.tensor(
                    clip_counts, dtype=torch.long
                )
            elif return_tensors == "np":
                audio_inputs["audio_num_clips"] = np.asarray(
                    clip_counts, dtype=np.int64
                )
            else:
                audio_inputs["audio_num_clips"] = clip_counts
        return BatchFeature(data={**encoded, **audio_inputs})

    def apply_transcription_request(
        self, audio, *, prompt=None, language_hint=None,
        audio_chunk_frames=None, return_tensors="pt",
    ):
        """Prepare one audio request in the released ASR input format.
        ``prompt`` starts with the base ASR instruction and may append topic
        and hotwords without a newline. A known language can be prefixed to
        the generated answer with ``language_hint='zh'`` or ``'en'``.
        """
        if prompt is None:
            prompt = _ASR_INSTRUCTION
        if not isinstance(prompt, str) or not prompt.startswith(_ASR_INSTRUCTION):
            raise ValueError("Prompt must start with the base ASR instruction")
        if "\n" in prompt or "\r" in prompt:
            raise ValueError("Prompt metadata must be contiguous, without newlines")
        if _PLACEHOLDER in prompt:
            raise ValueError("Prompt must not contain <AUDIO>; it is appended here")
        if audio_chunk_frames is not None:
            if (isinstance(audio_chunk_frames, bool)
                    or not isinstance(audio_chunk_frames, int)
                    or audio_chunk_frames < 1):
                raise ValueError("audio_chunk_frames must be a positive integer")
            extractor = self.feature_extractor
            original_config = extractor.audio_config
            extractor.audio_config = {**original_config, "chunk_frames": audio_chunk_frames}
            try:
                _ = extractor.spec
                inputs = self(
                    text=prompt + _PLACEHOLDER, audio=audio,
                    return_tensors=return_tensors,
                )
            finally:
                extractor.audio_config = original_config
        else:
            inputs = self(
                text=prompt + _PLACEHOLDER, audio=audio,
                return_tensors=return_tensors,
            )
        if language_hint is not None:
            label = {"zh": "Chinese\t", "en": "English\t"}.get(language_hint)
            if label is None:
                raise ValueError("language_hint must be 'zh', 'en', or None")
            if return_tensors != "pt":
                raise ValueError("language_hint requires return_tensors='pt'")
            label_ids = self.tokenizer(
                label, add_special_tokens=False, return_tensors="pt",
            )["input_ids"]
            if inputs.input_ids.shape[0] != 1:
                raise ValueError("language_hint requires exactly one audio request")
            inputs["input_ids"] = torch.cat((inputs.input_ids, label_ids), dim=1)
            if "attention_mask" in inputs:
                inputs["attention_mask"] = torch.cat(
                    (inputs.attention_mask, torch.ones_like(label_ids)), dim=1,
                )
        return inputs

    @property
    def model_input_names(self):
        return [
            "input_ids",
            "attention_mask",
            "audio_feats",
            "audio_feats_lengths",
            "audio_num_clips",
        ]


MingV3ASRProcessor.register_for_auto_class()

	# Copyright 2026 Ant Group. All rights reserved.
	# 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 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.
	"""Hugging Face processor for Ming V3 ASR transcription."""

	import re

	import numpy as np
	import torch
	from transformers.feature_extraction_utils import BatchFeature
	from transformers.processing_utils import ProcessingKwargs, ProcessorMixin

	from .audio_spec import MingAudioSpec # noqa: F401 - include transitive dependency in HF's local module cache
	from .feature_extraction_ming_v3_asr import MingV3ASRFeatureExtractor

	_AUDIO_FORMS = (
	("<audio>", "<audioPatch>", "</audio>"),
	("<\|audio_start\|>", "<\|audio_pad\|>", "<\|audio_end\|>"),
	)
	_PLACEHOLDER = "<AUDIO>"
	_ASR_INSTRUCTION = (
	"Please recognize the language of this speech and transcribe it, "
	"placing a `\t` between the identified language and the transcribe."
	)


	class MingV3ASRProcessorKwargs(ProcessingKwargs, total=False):
	_defaults = {
	"text_kwargs": {
	"padding": False,
	"padding_side": "left",
	"return_attention_mask": True,
	"return_token_type_ids": False,
	},
	}


	class MingV3ASRProcessor(ProcessorMixin):
	"""Expand audio tokens and combine an unchanged tokenizer with log-mel.

	``audio`` (or the compatible ``audios`` spelling) is a flat list of clips in
	prompt-placeholder order. ``audio_num_clips`` records each text request's
	clip count, including zero. Audio features and raw lengths are flattened
	independently of the text batch. Image/video inputs are not supported.

	Chat-template support is inherited from ``ProcessorMixin`` for tokenizer
	compatibility. ASR inference uses ``apply_transcription_request`` instead.
	"""

	feature_extractor_class = "AutoFeatureExtractor"
	tokenizer_class = "AutoTokenizer"
	valid_processor_kwargs = MingV3ASRProcessorKwargs

	def __init__(
	self,
	feature_extractor: MingV3ASRFeatureExtractor,
	tokenizer,
	chat_template=None,
	):
	template = tokenizer.chat_template if chat_template is None else chat_template
	super().__init__(feature_extractor, tokenizer, chat_template=template)
	vocabulary = tokenizer.get_vocab()
	self.audio_start_token, self.audio_token, self.audio_end_token = next(
	(
	form
	for form in _AUDIO_FORMS
	if all(token in vocabulary for token in form)
	),
	(None, None, None),
	)
	if self.audio_token is None:
	raise ValueError("tokenizer lacks Ming audio start/patch/end tokens")
	self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)

	@staticmethod
	def _normalize_prompt(text):
	for start, patch, end in _AUDIO_FORMS:
	text = re.sub(
	re.escape(start) + "(?:" + re.escape(patch) + ")+" + re.escape(end),
	_PLACEHOLDER,
	text,
	)
	return text

	def __call__(
	self, text=None, audio=None, , images=None, videos=None, audios=None, *kwargs
	):
	if images is not None or videos is not None:
	raise ValueError("MingV3ASRProcessor supports audio and text only")
	if audio is not None and audios is not None:
	raise ValueError("Pass either audio or audios, not both")
	clips = self.feature_extractor.as_clips(audio if audio is not None else audios)
	if text is None:
	if not clips:
	raise ValueError("Provide text or audio")
	text = " ".join([_PLACEHOLDER] * len(clips))
	texts = [text] if isinstance(text, str) else text
	if (
	not isinstance(texts, (list, tuple))
	or not texts
	or not all(isinstance(t, str) for t in texts)
	):
	raise ValueError("text must be a string or a nonempty list of strings")
	texts = [self._normalize_prompt(t) for t in texts]
	clip_counts = [t.count(_PLACEHOLDER) for t in texts]
	if sum(clip_counts) != len(clips):
	raise ValueError(
	f"Audio placeholder count {sum(clip_counts)} does not match {len(clips)} clips"
	)

	options = self._merge_kwargs(
	MingV3ASRProcessorKwargs,
	tokenizer_init_kwargs=self.tokenizer.init_kwargs,
	**kwargs,
	)
	text_options = options["text_kwargs"]
	audio_options = options["audio_kwargs"]
	return_tensors = text_options.get("return_tensors")
	audio_return_tensors = audio_options.get("return_tensors")
	if (
	return_tensors is not None
	and audio_return_tensors is not None
	and return_tensors != audio_return_tensors
	):
	raise ValueError("text and audio return_tensors must match")
	if return_tensors is None:
	return_tensors = audio_return_tensors
	if return_tensors is not None:
	text_options["return_tensors"] = return_tensors
	audio_options["return_tensors"] = return_tensors
	if text_options.get("return_overflowing_tokens"):
	raise ValueError("Audio token overflow windows are not supported")
	# Top-level max_length/truncation refer to text; audio is always chunked.
	# Explicit nested audio options retain extractor validation.
	for key in ("max_length", "truncation", "padding"):
	if key not in kwargs.get("audio_kwargs", {}):
	audio_options.pop(key, None)
	audio_inputs = self.feature_extractor(clips, **audio_options) if clips else {}
	lengths = audio_inputs.get("audio_feats_lengths", [])
	token_counts = [
	self.feature_extractor.spec.token_count(int(length)) for length in lengths
	]

	expected_counts = []
	clip_index = 0
	for index, count in enumerate(clip_counts):
	expected_counts.append(sum(token_counts[clip_index : clip_index + count]))
	for token_count in token_counts[clip_index : clip_index + count]:
	replacement = (
	self.audio_start_token
	+ self.audio_token * token_count
	+ self.audio_end_token
	)
	texts[index] = texts[index].replace(_PLACEHOLDER, replacement, 1)
	clip_index += count
	encoded = self.tokenizer(texts, **text_options)
	for ids, expected in zip(encoded["input_ids"], expected_counts):
	actual = sum(int(token) == self.audio_token_id for token in ids)
	if actual != expected:
	raise ValueError(
	"Audio token count changed during tokenization; do not truncate audio placeholders"
	)
	if clips:
	if return_tensors == "pt":
	audio_inputs["audio_num_clips"] = torch.tensor(
	clip_counts, dtype=torch.long
	)
	elif return_tensors == "np":
	audio_inputs["audio_num_clips"] = np.asarray(
	clip_counts, dtype=np.int64
	)
	else:
	audio_inputs["audio_num_clips"] = clip_counts
	return BatchFeature(data={encoded, audio_inputs})

	def apply_transcription_request(
	self, audio, *, prompt=None, language_hint=None,
	audio_chunk_frames=None, return_tensors="pt",
	):
	"""Prepare one audio request in the released ASR input format.

	``prompt`` starts with the base ASR instruction and may append topic
	and hotwords without a newline. A known language can be prefixed to
	the generated answer with ``language_hint='zh'`` or ``'en'``.
	"""
	if prompt is None:
	prompt = _ASR_INSTRUCTION
	if not isinstance(prompt, str) or not prompt.startswith(_ASR_INSTRUCTION):
	raise ValueError("Prompt must start with the base ASR instruction")
	if "\n" in prompt or "\r" in prompt:
	raise ValueError("Prompt metadata must be contiguous, without newlines")
	if _PLACEHOLDER in prompt:
	raise ValueError("Prompt must not contain <AUDIO>; it is appended here")
	if audio_chunk_frames is not None:
	if (isinstance(audio_chunk_frames, bool)
	or not isinstance(audio_chunk_frames, int)
	or audio_chunk_frames < 1):
	raise ValueError("audio_chunk_frames must be a positive integer")
	extractor = self.feature_extractor
	original_config = extractor.audio_config
	extractor.audio_config = {**original_config, "chunk_frames": audio_chunk_frames}
	try:
	_ = extractor.spec
	inputs = self(
	text=prompt + _PLACEHOLDER, audio=audio,
	return_tensors=return_tensors,
	)
	finally:
	extractor.audio_config = original_config
	else:
	inputs = self(
	text=prompt + _PLACEHOLDER, audio=audio,
	return_tensors=return_tensors,
	)
	if language_hint is not None:
	label = {"zh": "Chinese\t", "en": "English\t"}.get(language_hint)
	if label is None:
	raise ValueError("language_hint must be 'zh', 'en', or None")
	if return_tensors != "pt":
	raise ValueError("language_hint requires return_tensors='pt'")
	label_ids = self.tokenizer(
	label, add_special_tokens=False, return_tensors="pt",
	)["input_ids"]
	if inputs.input_ids.shape[0] != 1:
	raise ValueError("language_hint requires exactly one audio request")
	inputs["input_ids"] = torch.cat((inputs.input_ids, label_ids), dim=1)
	if "attention_mask" in inputs:
	inputs["attention_mask"] = torch.cat(
	(inputs.attention_mask, torch.ones_like(label_ids)), dim=1,
	)
	return inputs

	@property
	def model_input_names(self):
	return [
	"input_ids",
	"attention_mask",
	"audio_feats",
	"audio_feats_lengths",
	"audio_num_clips",
	]


	MingV3ASRProcessor.register_for_auto_class()