# coding=utf-8 from __future__ import annotations import base64 import io import json import os from typing import Any, Dict, List, Optional, Union import numpy as np import torch import librosa import soundfile as sf # Explicitly import soundfile to handle BytesIO. from transformers import AutoTokenizer, WhisperFeatureExtractor from transformers.feature_extraction_utils import BatchFeature from transformers.processing_utils import ProcessorMixin from transformers.utils import logging logger = logging.get_logger(__name__) _AUDIO_MARKER = "<>" def _normalize_dtype_name(name: str) -> str: name = name.strip().lower() alias = { "fp16": "float16", "float16": "float16", "half": "float16", "bf16": "bfloat16", "bfloat16": "bfloat16", "fp32": "float32", "float32": "float32", "float": "float32", } return alias.get(name, name) def _resolve_torch_dtype(x: Any, default: str = "float32") -> torch.dtype: if isinstance(x, torch.dtype): return x if x is None: x = default if isinstance(x, str): name = _normalize_dtype_name(x) if not hasattr(torch, name): raise ValueError(f"Unknown torch dtype string: {x} (normalized: {name})") return getattr(torch, name) raise TypeError(f"audio_dtype/audio_torch_dtype must be str or torch.dtype or None, got {type(x)}") class ArkasrProcessor(ProcessorMixin): attributes = ["feature_extractor", "tokenizer"] valid_kwargs = ["merge_factor", "audio_token", "audio_dtype"] feature_extractor_class = ("WhisperFeatureExtractor", "SequenceFeatureExtractor") tokenizer_class = ("PreTrainedTokenizerFast", "PreTrainedTokenizer") def __init__( self, feature_extractor, tokenizer, merge_factor: int = 4, audio_token: str = "<|audio|>", audio_dtype: str = "float32", **kwargs, ): super().__init__(feature_extractor, tokenizer) self.merge_factor = int(merge_factor) self.audio_token = str(audio_token) self.audio_dtype = str(audio_dtype) self.bos_audio_token = "<|begin_of_audio|>" self.eos_audio_token = "<|end_of_audio|>" self.user_token = "<|user|>" self.assistant_token = "<|assistant|>" @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "ArkasrProcessor": trust_remote_code = bool(kwargs.pop("trust_remote_code", False)) passthrough_keys = {"cache_dir", "force_download", "local_files_only", "token", "revision", "subfolder"} shared_kwargs = {k: kwargs[k] for k in list(kwargs.keys()) if k in passthrough_keys} merge_factor = 4 audio_token = "<|audio|>" audio_dtype = "float32" tokenizer_cfg: Dict[str, Any] = {} feat_cfg: Dict[str, Any] = {} proc_cfg_path = os.path.join(pretrained_model_name_or_path, "processor_config.json") if os.path.isfile(proc_cfg_path): with open(proc_cfg_path, "r", encoding="utf-8") as f: proc_cfg = json.load(f) merge_factor = int(proc_cfg.get("merge_factor", merge_factor)) audio_token = str(proc_cfg.get("audio_token", audio_token)) audio_dtype = str(proc_cfg.get("audio_dtype", audio_dtype)) tokenizer_cfg = proc_cfg.get("tokenizer_config", {}) or {} feat_cfg = proc_cfg.get("feature_extractor_config", {}) or {} feature_extractor = WhisperFeatureExtractor.from_pretrained(pretrained_model_name_or_path, **shared_kwargs) for k, v in feat_cfg.items(): if hasattr(feature_extractor, k): try: setattr(feature_extractor, k, v) except Exception: pass tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path, use_fast=True, trust_remote_code=trust_remote_code, **shared_kwargs ) for k, v in tokenizer_cfg.items(): if hasattr(tokenizer, k): try: setattr(tokenizer, k, v) except Exception: pass return cls( feature_extractor=feature_extractor, tokenizer=tokenizer, merge_factor=merge_factor, audio_token=audio_token, audio_dtype=audio_dtype, ) # ========================= # audio helpers (Modified) # ========================= def _load_audio_file(self, path: str, sampling_rate: int = 16000, offset: float = 0.0, duration: Optional[float] = None) -> np.ndarray: # librosa.load supports offset and duration. # offset: start reading after this time (in seconds) # duration: only load up to this much audio (in seconds) audio_array, _ = librosa.load(path, sr=int(sampling_rate), mono=True, offset=offset, duration=duration) return np.asarray(audio_array, dtype=np.float32) def _strip_data_url_prefix(self, b64: str) -> str: if "," in b64 and b64[:30].lower().startswith("data:"): return b64.split(",", 1)[1] return b64 def _load_audio_base64(self, b64: str, sampling_rate: int = 16000, offset: float = 0.0, duration: Optional[float] = None) -> np.ndarray: b64 = self._strip_data_url_prefix(b64) raw = base64.b64decode(b64) bio = io.BytesIO(raw) # librosa also supports offset and duration when loading from BytesIO. try: wav, _sr = librosa.load(bio, sr=int(sampling_rate), mono=True, offset=offset, duration=duration) return np.asarray(wav, dtype=np.float32) except Exception as e: # Fallback path: manual slicing, which is slower. try: bio.seek(0) data, sr = sf.read(bio, dtype="float32", always_2d=True) wav = data.mean(axis=1) if int(sr) != int(sampling_rate): wav = librosa.resample(wav, orig_sr=int(sr), target_sr=int(sampling_rate)) start_sample = int(offset * sampling_rate) end_sample = None if duration is not None: end_sample = start_sample + int(duration * sampling_rate) return np.asarray(wav[start_sample:end_sample], dtype=np.float32) except Exception as e2: raise ValueError("Failed to decode base64 audio.") from e2 def calculate_audio_token_count(self, mel_frames: int) -> int: downsampled = (int(mel_frames) + 1) // 2 merged = downsampled // max(self.merge_factor, 1) return max(int(merged), 1) def _build_templates_and_audios( self, conversations: List[List[dict]], sampling_rate: int, add_generation_prompt: bool, ) -> tuple[List[str], List[np.ndarray], List[int]]: prompts_template: List[str] = [] audios_raw: List[np.ndarray] = [] prompt_audio_counts: List[int] = [] for conv in conversations: conv_str = "" last_role = None audio_count_this_conv = 0 for msg in conv: role = msg["role"] last_role = role content = msg["content"] if role == "user": conv_str += f"{self.user_token}" elif role == "assistant": conv_str += f"{self.assistant_token}" else: conv_str += f"<|{role}|>" if isinstance(content, str): conv_str += f"{content}" elif isinstance(content, list): for part in content: ptype = part.get("type") if ptype == "audio": # ------------------------------------------------------------ # Parse begin_time and end_time when present. # ------------------------------------------------------------ begin_time = part.get("begin_time", -1) end_time = part.get("end_time", -1) offset = 0.0 duration = None # Apply slicing only when begin_time is valid and non-negative. if begin_time is not None and begin_time >= 0: offset = float(begin_time) if end_time is not None and end_time > begin_time: duration = float(end_time) - float(begin_time) audio_raw_this = None if "array" in part: arr = part["array"] if isinstance(arr, torch.Tensor): arr = arr.detach().cpu().numpy() full_arr = np.asarray(arr, dtype=np.float32).reshape(-1) # Slice the in-memory audio array. start_idx = int(offset * sampling_rate) end_idx = None if duration is not None: end_idx = start_idx + int(duration * sampling_rate) audio_raw_this = full_arr[start_idx:end_idx] elif "path" in part: audio_raw_this = self._load_audio_file( part["path"], sampling_rate=sampling_rate, offset=offset, duration=duration ) elif "base64" in part: audio_raw_this = self._load_audio_base64( part["base64"], sampling_rate=sampling_rate, offset=offset, duration=duration ) else: raise ValueError("Audio part must contain 'path' or 'array' or 'base64'.") audios_raw.append(audio_raw_this) audio_count_this_conv += 1 conv_str += f"{self.bos_audio_token}{_AUDIO_MARKER}{self.eos_audio_token}" elif ptype == "text": conv_str += f"{part.get('text', '')}" else: raise ValueError(f"Unknown content part type: {ptype}") else: raise ValueError(f"Unsupported message content type: {type(content)}") if add_generation_prompt: if last_role == "user": conv_str += f"{self.assistant_token}" elif last_role == "assistant": conv_str += f"{self.user_token}" else: conv_str += f"{self.assistant_token}" prompts_template.append(conv_str) prompt_audio_counts.append(audio_count_this_conv) return prompts_template, audios_raw, prompt_audio_counts def _calculate_audio_token_counts_per_sample( self, audios_raw: List[np.ndarray], sampling_rate: int, audio_max_length: Optional[int], audio_pad_to_multiple_of: Optional[int], ) -> List[int]: del sampling_rate, audio_pad_to_multiple_of hop_length = int(getattr(self.feature_extractor, "hop_length", 160)) max_audio_samples = int(audio_max_length) if audio_max_length is not None else None token_counts: List[int] = [] for audio_raw in audios_raw: audio_np = np.asarray(audio_raw, dtype=np.float32).reshape(-1) effective_len = int(audio_np.shape[0]) if max_audio_samples is not None: effective_len = min(effective_len, max_audio_samples) mel_frames = effective_len // max(hop_length, 1) token_counts.append(self.calculate_audio_token_count(int(mel_frames))) return token_counts # ========================= # apply_chat_template # ========================= def apply_chat_template( self, conversation: Union[List[dict], List[List[dict]]], chat_template: Optional[str] = None, add_generation_prompt: bool = True, **kwargs, ) -> Union[BatchFeature, str, List[str]]: if chat_template is not None: logger.warning("chat_template argument is ignored.") tokenize = kwargs.pop("tokenize", True) return_tensors = kwargs.pop("return_tensors", "pt") kwargs.pop("return_dict", None) audio_torch_dtype = kwargs.pop("audio_torch_dtype", None) audio_dtype_override = kwargs.pop("audio_dtype", None) dtype_source = audio_torch_dtype if audio_torch_dtype is not None else audio_dtype_override target_dtype = _resolve_torch_dtype(dtype_source, default=getattr(self, "audio_dtype", "float32")) text_kwargs = dict(kwargs.pop("text_kwargs", {}) or {}) for k in ("padding", "truncation", "max_length", "add_special_tokens"): if k in kwargs and k not in text_kwargs: text_kwargs[k] = kwargs.pop(k) sampling_rate = int(kwargs.pop("sampling_rate", 16000)) audio_padding = kwargs.pop("audio_padding", "longest") audio_max_length = kwargs.pop("audio_max_length", None) audio_pad_to_multiple_of = kwargs.pop("audio_pad_to_multiple_of", None) if kwargs: logger.warning(f"Ignored unused kwargs: {list(kwargs.keys())}") if isinstance(conversation, list) and conversation and isinstance(conversation[0], dict): conversations = [conversation] is_single = True else: conversations = conversation is_single = False prompt_templates, audios_raw, prompt_audio_counts = self._build_templates_and_audios( conversations=conversations, sampling_rate=sampling_rate, add_generation_prompt=add_generation_prompt, ) input_features = None audio_token_counts: List[int] = [] if len(audios_raw) > 0: feat = self.feature_extractor( audios_raw, sampling_rate=sampling_rate, return_tensors="np", return_attention_mask=False, padding=audio_padding, max_length=audio_max_length, pad_to_multiple_of=audio_pad_to_multiple_of, ) input_features = feat["input_features"] if not isinstance(input_features, np.ndarray): input_features = np.asarray(input_features) audio_token_counts = self._calculate_audio_token_counts_per_sample( audios_raw=audios_raw, sampling_rate=sampling_rate, audio_max_length=audio_max_length, audio_pad_to_multiple_of=audio_pad_to_multiple_of, ) prompts: List[str] = [] audio_idx = 0 for prompt_template, audio_count in zip(prompt_templates, prompt_audio_counts): prompt = prompt_template for _ in range(audio_count): if audio_idx >= len(audio_token_counts): raise ValueError("Audio token count mismatch while building prompts.") audio_tokens_str = "".join([self.audio_token] * audio_token_counts[audio_idx]) prompt = prompt.replace(_AUDIO_MARKER, audio_tokens_str, 1) audio_idx += 1 if _AUDIO_MARKER in prompt: raise ValueError("Unresolved audio marker remained in prompt.") prompts.append(prompt) if audio_idx != len(audio_token_counts): raise ValueError("Unused audio token counts remained after prompt construction.") if not tokenize: return prompts[0] if is_single else prompts text_kwargs.setdefault("padding", "longest") text_kwargs.setdefault("add_special_tokens", False) text_kwargs["return_tensors"] = return_tensors enc = self.tokenizer(prompts, **text_kwargs) data: Dict[str, Any] = dict(enc) if input_features is not None: data["audios"] = torch.tensor(input_features, dtype=target_dtype) return BatchFeature(data=data, tensor_type=return_tensors) # ... (The remaining batch_decode, decode, __call__, and model_input_names stay unchanged.) ... def batch_decode(self, *args, **kwargs): return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): return self.tokenizer.decode(*args, **kwargs) def __call__( self, text: Union[str, List[str]], audios: Union[np.ndarray, torch.Tensor, List[Union[np.ndarray, torch.Tensor]]], sampling_rate: int = 16000, return_tensors: str = "pt", **tokenizer_kwargs, ) -> BatchFeature: # Simplified implementation that skips time slicing because the caller passes raw audio arrays directly. audios_list = [] def flatten_audios(obj): if isinstance(obj, (list, tuple)): if len(obj) > 0 and isinstance(obj[0], (float, int)): audios_list.append(obj) else: for item in obj: flatten_audios(item) elif isinstance(obj, (np.ndarray, torch.Tensor)): audios_list.append(obj) flatten_audios(audios) audios_np: List[np.ndarray] = [] for a in audios_list: if isinstance(a, torch.Tensor): a = a.detach().cpu().numpy() a = np.asarray(a, dtype=np.float32).reshape(-1) audios_np.append(a) input_features = None if audios_np: feat = self.feature_extractor(audios_np, sampling_rate=int(sampling_rate), return_tensors="np", return_attention_mask=False, padding="longest") input_features = feat["input_features"] if not isinstance(input_features, np.ndarray): input_features = np.asarray(input_features) tokenizer_kwargs = dict(tokenizer_kwargs or {}) tokenizer_kwargs.setdefault("padding", "longest") tokenizer_kwargs.setdefault("add_special_tokens", False) tokenizer_kwargs["return_tensors"] = return_tensors enc = self.tokenizer(text, **tokenizer_kwargs) data: Dict[str, Any] = dict(enc) if input_features is not None: data["audios"] = torch.tensor(input_features, dtype=_resolve_torch_dtype(getattr(self, "audio_dtype", "float32"))) return BatchFeature(data=data, tensor_type=return_tensors) @property def model_input_names(self): return ["input_ids", "attention_mask", "audios"]