Upload GPA v1.5 model package

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+.DS_Store

README.md

@@ -0,0 +1,3 @@
+---
+license: apache-2.0
+---

chat_template.jinja

@@ -0,0 +1,15 @@
+{% for message in messages %}
+    {% if message.role == "system" %}
+<|system|>
+{{ message.content }}
+    {% elif message.role == "user" %}
+<|user|>
+{{ message.content }}
+    {% elif message.role == "assistant" %}
+<|assistant|>
+{{ message.content }}
+    {% endif %}
+{% endfor %}
+{% if add_generation_prompt %}
+<|assistant|>
+{% endif %}

config.json

@@ -0,0 +1,112 @@
+{
+  "adapter_type": "mlp",
+  "architectures": [
+    "ArkasrForConditionalGeneration"
+  ],
+  "attention_dropout": 0.0,
+  "audio_token_id": 151663,
+  "auto_map": {
+    "AutoConfig": "configuration_arkasr.ArkasrConfig",
+    "AutoModelForCausalLM": "modeling_arkasr.ArkasrForConditionalGeneration"
+  },
+  "dtype": "float32",
+  "eos_token_id": 151665,
+  "hidden_act": "silu",
+  "hidden_size": 896,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_whisper_length": 1500,
+  "max_window_layers": 24,
+  "merge_factor": 4,
+  "mlp_adapter_act": "gelu",
+  "model_type": "arkasr",
+  "num_attention_heads": 14,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 2,
+  "pad_token_id": 151643,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "spec_aug": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.3",
+  "use_cache": false,
+  "use_mrope": false,
+  "use_rope": true,
+  "use_sliding_window": false,
+  "vocab_size": 163958,
+  "whisper_config": {
+    "activation_dropout": 0.0,
+    "activation_function": "gelu",
+    "apply_spec_augment": false,
+    "architectures": [
+      "WhisperForConditionalGeneration"
+    ],
+    "attention_dropout": 0.0,
+    "begin_suppress_tokens": [
+      220,
+      50257
+    ],
+    "bos_token_id": 50257,
+    "classifier_proj_size": 256,
+    "d_model": 1280,
+    "decoder_attention_heads": 20,
+    "decoder_ffn_dim": 5120,
+    "decoder_layerdrop": 0.0,
+    "decoder_layers": 32,
+    "decoder_start_token_id": 50258,
+    "dropout": 0.0,
+    "dtype": "bfloat16",
+    "encoder_attention_heads": 20,
+    "encoder_ffn_dim": 5120,
+    "encoder_layerdrop": 0.0,
+    "encoder_layers": 32,
+    "eos_token_id": 50257,
+    "init_std": 0.02,
+    "mask_feature_length": 10,
+    "mask_feature_min_masks": 0,
+    "mask_feature_prob": 0.0,
+    "mask_time_length": 10,
+    "mask_time_min_masks": 2,
+    "mask_time_prob": 0.05,
+    "max_length": 448,
+    "max_source_positions": 1500,
+    "max_target_positions": 448,
+    "median_filter_width": 7,
+    "model_type": "whisper",
+    "num_hidden_layers": 32,
+    "num_mel_bins": 128,
+    "scale_embedding": false,
+    "use_cache": true,
+    "use_weighted_layer_sum": false,
+    "vocab_size": 51866
+  }
+}

configuration_arkasr.py

@@ -0,0 +1,47 @@
+from typing import Any, Dict, Optional, Union
+from transformers import Qwen2Config, WhisperConfig
+
+
+class ArkasrConfig(Qwen2Config):
+    model_type = "arkasr"
+    is_composition = True
+
+    def __init__(
+        self,
+        whisper_config: Optional[Union[Dict[str, Any], WhisperConfig]] = None,
+        adapter_type: str = "mlp",
+        merge_factor: int = 4,
+        spec_aug: bool = False,
+        use_rope: bool = True,
+        max_whisper_length: int = 1500,
+        mlp_adapter_act: str = "gelu",
+        **kwargs,  # All Qwen2Config parameters are forwarded from here.
+    ):
+        # === 1️⃣ Key point: initialize Qwen2Config (LM section) ===
+        # This consumes fields such as:
+        # vocab_size / hidden_size / num_hidden_layers / rope_scaling / ...
+        super().__init__(**kwargs)
+
+        # === 2️⃣ Whisper sub-config ===
+        if isinstance(whisper_config, dict):
+            self.whisper_config = WhisperConfig(**whisper_config)
+        elif isinstance(whisper_config, WhisperConfig):
+            self.whisper_config = whisper_config
+        else:
+            self.whisper_config = WhisperConfig()
+
+        # === 3️⃣ ArkASR-specific parameters ===
+        self.adapter_type = adapter_type
+        self.merge_factor = int(merge_factor)
+        self.spec_aug = bool(spec_aug)
+        self.use_rope = bool(use_rope)
+        self.max_whisper_length = int(max_whisper_length)
+        self.mlp_adapter_act = mlp_adapter_act
+
+    def to_dict(self):
+        output = super().to_dict()
+        output["whisper_config"] = self.whisper_config.to_dict()
+        return output
+
+
+__all__ = ["ArkasrConfig"]

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 151665,
+  "pad_token_id": 151665,
+  "transformers_version": "4.57.3"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a7d69140f7b75d815e2136a83fc307dc353f6d5ecd2134a1187d6948ae6aa37
+size 2305207856

modeling_arkasr.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+
+from typing import Optional, List, Tuple, Union, Dict
+
+import torch
+from torch import Tensor, nn
+from transformers import Qwen2ForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_arkasr import ArkasrConfig
+from .modeling_audio import WhisperSpecialEncoder
+
+
+class AudioMLPAdapter(nn.Module):
+    def __init__(self, config: ArkasrConfig):
+        super().__init__()
+        whisper_config = config.whisper_config
+        self.merge_factor = int(config.merge_factor)
+
+        # Audio encoder
+        self.whisper = WhisperSpecialEncoder(
+            whisper_config,
+            use_rope=getattr(config, "use_rope", False),
+        )
+        # Disable Whisper's built-in LayerNorm.
+        self.whisper.layer_norm = nn.Identity()
+        self.layer_norm = nn.LayerNorm(whisper_config.hidden_size)
+
+        act_fn_map = {
+            "gelu": nn.GELU(),
+            "relu": nn.ReLU(),
+            "selu": nn.SELU(),
+        }
+        act = act_fn_map.get(getattr(config, "mlp_adapter_act", "gelu"), nn.GELU())
+
+        input_dim = whisper_config.hidden_size * self.merge_factor
+        output_dim = config.hidden_size
+
+        self.adapting = nn.Sequential(
+            nn.Linear(input_dim, output_dim * 2),
+            act,
+            nn.Linear(output_dim * 2, output_dim),
+        )
+
+    def forward(self, audios: Tensor) -> Tensor:
+        """
+        Args:
+            audios: (B, mel, T) or (B, raw_len), depending on WhisperSpecialEncoder.
+        Returns:
+            adapted_features: (B, Seq_Audio, LLM_Hidden_Dim)
+        """
+        bsz = audios.size(0)
+
+        encoded = self.whisper(audios)[0]  # (B, T, D)
+        encoded = self.layer_norm(encoded)
+
+        seq_len = encoded.size(1)
+        if seq_len % self.merge_factor != 0:
+            target_len = (seq_len // self.merge_factor) * self.merge_factor
+            if target_len <= 0:
+                # Guard for extremely short audio: pad to merge_factor.
+                target_len = self.merge_factor
+                if seq_len < target_len:
+                    pad_len = target_len - seq_len
+                    pad = encoded.new_zeros((bsz, pad_len, encoded.size(-1)))
+                    encoded = torch.cat([encoded, pad], dim=1)
+            else:
+                encoded = encoded[:, :target_len, :]
+
+        encoded = encoded.reshape(bsz, -1, encoded.size(-1) * self.merge_factor)
+        adapted = self.adapting(encoded)  # (B, T/k, hidden)
+        return adapted
+
+
+class ArkasrForConditionalGeneration(Qwen2ForCausalLM):
+    config_class = ArkasrConfig
+    _no_split_modules = ["WhisperSpecialEncoder"]
+
+    def __init__(self, config: ArkasrConfig):
+        super().__init__(config)
+        self.audio_encoder = AudioMLPAdapter(config)
+
+        self.audio_token_id = getattr(config, "audio_token_id", None)
+        if self.audio_token_id is None:
+            raise ValueError("`audio_token_id` must be defined in config.")
+
+    @staticmethod
+    def _cache_seq_len(past_key_values) -> int:
+        if past_key_values is None:
+            return 0
+        if hasattr(past_key_values, "get_seq_length"):
+            try:
+                return int(past_key_values.get_seq_length())
+            except Exception:
+                return 0
+        try:
+            return int(past_key_values[0][0].shape[-2])
+        except Exception:
+            return 0
+
+    def _inject_audio_embeddings_batch_encode_then_loop_scatter(
+        self,
+        input_ids: torch.LongTensor,          # (B, S)
+        inputs_embeds: torch.FloatTensor,     # (B, S, H)
+        audios: Tensor,                       # (B, ...)
+    ) -> torch.FloatTensor:
+        """
+                First run one batched audio encoding pass for samples that contain audio tokens,
+                then scatter each sample's audio features back into inputs_embeds at the
+                corresponding audio_token positions.
+
+                Benefits:
+                - The encoder runs only once.
+                - Scatter is performed per sample, so features cannot drift across samples.
+                - Rows without audio_token are skipped directly, which keeps TTS-only rows unaffected.
+
+                Constraint:
+                - The number of audio tokens n_i in each sample should align with Sa from the
+                    audio encoder output. If they do not align, this path truncates or zero-pads
+                    to n_i instead of raising an error.
+        """
+        B, S = input_ids.shape
+        H = inputs_embeds.size(-1)
+        device = inputs_embeds.device
+        dtype = inputs_embeds.dtype
+
+        # Find the samples that require audio injection.
+        mask = (input_ids == self.audio_token_id)  # (B, S)
+        per_counts = mask.sum(dim=1)               # (B,)
+        need_idx = (per_counts > 0).nonzero(as_tuple=False).squeeze(1)  # (K,)
+
+        if need_idx.numel() == 0:
+            return inputs_embeds
+
+        # Encode only the subset of audio that needs injection. (K, ...)
+        audios_sub = audios.index_select(0, need_idx)
+        feats_sub = self.audio_encoder(audios_sub)  # (K, Sa, H)
+
+        # Scatter back per sample; the write-back itself is negligible.
+        feats_sub = feats_sub.to(device=device, dtype=dtype)
+        Sa = feats_sub.size(1)
+
+        # Inject one sample at a time.
+        for k in range(need_idx.numel()):
+            i = int(need_idx[k].item())
+            n_i = int(per_counts[i].item())
+            if n_i <= 0:
+                continue
+
+            feat_i = feats_sub[k]  # (Sa, H)
+
+            # Align to the number of audio tokens n_i for this sample.
+            if Sa < n_i:
+                pad = feat_i.new_zeros((n_i - Sa, H))
+                feat_i_use = torch.cat([feat_i, pad], dim=0)
+            elif Sa > n_i:
+                feat_i_use = feat_i[:n_i]
+            else:
+                feat_i_use = feat_i
+
+            pos_i = mask[i].nonzero(as_tuple=False).squeeze(1)  # (n_i,)
+            # Write features back into embeddings.
+            inputs_embeds[i, pos_i, :] = feat_i_use
+
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        audios: Optional[Tensor] = None,
+        attention_mask: Optional[Tensor] = None,
+        position_ids: Optional[Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("Either `input_ids` or `inputs_embeds` must be provided.")
+            inputs_embeds = self.model.embed_tokens(input_ids)
+
+        # Inject only on the first step (past_len == 0) to avoid re-encoding during generation.
+        past_len = self._cache_seq_len(past_key_values)
+        if audios is not None and input_ids is not None and past_len == 0:
+            inputs_embeds = self._inject_audio_embeddings_batch_encode_then_loop_scatter(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                audios=audios,
+            )
+
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        hidden_states = outputs[0]
+
+        # Restrict logits computation when possible to avoid redundant lm_head work.
+        if isinstance(logits_to_keep, int) and logits_to_keep > 0:
+            hidden_for_logits = hidden_states[:, -logits_to_keep:, :]
+        elif isinstance(logits_to_keep, torch.Tensor):
+            hidden_for_logits = hidden_states[:, logits_to_keep, :]
+        else:
+            hidden_for_logits = hidden_states
+
+        logits = self.lm_head(hidden_for_logits)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        **kwargs,
+    ):
+        past_len = self._cache_seq_len(past_key_values)
+        if past_len > 0:
+            input_ids = input_ids[:, -1:]
+
+        model_inputs = {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "attention_mask": attention_mask,
+            # Pass audios through. Injection happens only when past_len == 0 in forward,
+            # so later generation steps do not re-encode.
+            "audios": kwargs.get("audios", None),
+        }
+
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs["inputs_embeds"] = inputs_embeds
+            del model_inputs["input_ids"]
+
+        return model_inputs
+
+
+__all__ = ["ArkasrForConditionalGeneration", "AudioMLPAdapter"]

modeling_audio.py

@@ -0,0 +1,293 @@
+from typing import Any, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+from torch.nn.functional import scaled_dot_product_attention
+from transformers import WhisperConfig
+from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
+from transformers.models.whisper.modeling_whisper import WhisperEncoder, WhisperEncoderLayer
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+# ==========================================
+# 1. Core Rotary Embedding components
+# ==========================================
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, rope_ratio=1):
+        super().__init__()
+        self.dim = dim
+        self.rope_ratio = rope_ratio
+
+    @torch.no_grad()
+    def get_emb(self, seq_len: int, dtype: torch.dtype, device: torch.device, base: int = 10000):
+        """Generate the cached RoPE table."""
+        base = base * self.rope_ratio
+        # Compute the theta frequencies.
+        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.float, device=device) / self.dim))
+        
+        # Build the position indices.
+        t = torch.arange(seq_len, device=device, dtype=torch.float)
+        freqs = torch.outer(t, inv_freq) # [seq_len, dim/2]
+        
+        # Construct the cos/sin cache.
+        # Shape: [seq_len, dim/2, 2]
+        emb = torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1)
+        
+        if dtype in (torch.float16, torch.bfloat16):
+            emb = emb.to(dtype)
+        return emb
+
+def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
+    """
+    x: [batch, num_heads, seq_len, head_dim]
+    rope_cache: [1, seq_len, dim/2, 2]
+    """
+    b, nh, sq, hd = x.shape
+    rot_dim = rope_cache.shape[-2] * 2
+    
+    # Split x into the rotated and pass-through portions.
+    x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
+    
+    # Reshape x_rot to match rope_cache: [b, nh, sq, rot_dim/2, 2]
+    x_shaped = x_rot.reshape(b, nh, sq, rot_dim // 2, 2)
+    
+    # Apply the complex rotation: (a+bi)(c+di) = (ac-bd) + (ad+bc)i
+    cos = rope_cache[..., 0] # [1, sq, rot_dim/2]
+    sin = rope_cache[..., 1] # [1, sq, rot_dim/2]
+    
+    # Add the head dimension.
+    cos = cos.unsqueeze(1) # [1, 1, sq, rot_dim/2]
+    sin = sin.unsqueeze(1) # [1, 1, sq, rot_dim/2]
+    
+    x_out = torch.stack([
+        x_shaped[..., 0] * cos - x_shaped[..., 1] * sin,
+        x_shaped[..., 1] * cos + x_shaped[..., 0] * sin
+    ], dim=-1)
+    
+    x_out = x_out.flatten(3) # Merge the final two dimensions into rot_dim.
+    return torch.cat([x_out, x_pass], dim=-1)
+
+# ==========================================
+# 2. RoPE attention built on SDPA
+# ==========================================
+
+class WhisperRoPESdpaAttention(nn.Module):
+    """
+    Replace WhisperFlashAttention2 with PyTorch's native scaled_dot_product_attention.
+    """
+    def __init__(self, config: WhisperConfig, embed_dim: int, num_heads: int, dropout: float = 0.0):
+        super().__init__()
+        self.config = config
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        
+        # Standard Whisper projection layers.
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        
+        self.is_causal = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], None]:
+        
+        bsz, q_len, _ = hidden_states.size()
+
+        # 1. Project to queries, keys, and values.
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # 2. Reshape to [batch, heads, seq, dim] and keep memory contiguous.
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+        key_states = key_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+        value_states = value_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+        # 3. Apply RoPE.
+        if rotary_pos_emb is not None:
+            query_states = apply_rotary_pos_emb(query_states, rotary_pos_emb)
+            key_states = apply_rotary_pos_emb(key_states, rotary_pos_emb)
+
+        # 4. Align dtypes to avoid mismatches introduced by fp32 LayerNorm.
+        target_dtype = self.q_proj.weight.dtype
+        query_states = query_states.to(target_dtype)
+        key_states = key_states.to(target_dtype)
+        value_states = value_states.to(target_dtype)
+
+        # 5. Run SDPA. Do not apply manual scaling; SDPA handles it internally.
+        # If a 4D attention_mask is provided, SDPA applies it correctly.
+        attn_output = scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=self.is_causal,
+        )
+
+        # 6. Restore shape and apply the output projection.
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None, None
+
+# ==========================================
+# 3. Wrapped encoder layer and encoder
+# ==========================================
+
+class WhisperSpecialEncoderLayer(WhisperEncoderLayer):
+    def __init__(self, config: WhisperConfig):
+        super().__init__(config)
+        # Replace self-attention with the RoPE + SDPA implementation.
+        self.self_attn = WhisperRoPESdpaAttention(
+            config=config,
+            embed_dim=self.embed_dim,
+            num_heads=config.encoder_attention_heads,
+            dropout=config.attention_dropout,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Any]:
+        
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        
+        hidden_states, attn_weights, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            layer_head_mask=layer_head_mask,
+            output_attentions=output_attentions,
+            rotary_pos_emb=rotary_pos_emb,
+        )
+        
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        if hidden_states.dtype == torch.float16:
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        return (hidden_states, None) # Keep the tuple length aligned with the Whisper interface.
+
+class WhisperSpecialEncoder(WhisperEncoder):
+    def __init__(self, config: WhisperConfig, use_rope=True, rope_ratio=1):
+        super().__init__(config)
+        self.use_rope = use_rope
+        # Override the parent layer stack.
+        self.layers = nn.ModuleList(
+            [WhisperSpecialEncoderLayer(config) for _ in range(config.encoder_layers)]
+        )
+        
+        if use_rope:
+            # Compute the RoPE dimension, typically a subset of head_dim.
+            head_dim = config.d_model // config.encoder_attention_heads
+            self.rotary_embedding = RotaryEmbedding(head_dim // 2, rope_ratio)
+
+    def forward(
+        self,
+        input_features,
+        attention_mask=None,
+        head_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        position_ids=None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Whisper convolutional feature extraction.
+        inputs_embeds = nn.functional.gelu(self.conv1(input_features))
+        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
+        inputs_embeds = inputs_embeds.permute(0, 2, 1) # [B, T_down, D]
+
+        if self.use_rope:
+            # Build the rotary embedding cache.
+            rotary_embs = self.rotary_embedding.get_emb(
+                seq_len=inputs_embeds.shape[1],
+                dtype=inputs_embeds.dtype,
+                device=inputs_embeds.device
+            )
+            # Reshape to [1, seq_len, dim/2, 2] for broadcasting.
+            rotary_embs = rotary_embs.unsqueeze(0)
+            hidden_states = inputs_embeds 
+        else:
+            rotary_embs = None
+            # Fall back to absolute positional embeddings.
+            embed_pos = self.embed_positions.weight[:inputs_embeds.shape[1]]
+            hidden_states = inputs_embeds + embed_pos
+
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    encoder_layer.__call__,
+                    hidden_states,
+                    None, # attention_mask
+                    (head_mask[idx] if head_mask is not None else None),
+                    output_attentions,
+                    rotary_embs,
+                    position_ids,
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    attention_mask=None,
+                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
+                    output_attentions=output_attentions,
+                    rotary_pos_emb=rotary_embs,
+                    position_ids=position_ids,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[2],)
+
+        hidden_states = self.layer_norm(hidden_states)
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+            
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_states,
+            attentions=all_attentions,
+        )

preprocessor_config.json

@@ -0,0 +1,13 @@
+{
+  "chunk_length": 30,
+  "feature_extractor_type": "WhisperFeatureExtractor",
+  "feature_size": 128,
+  "hop_length": 160,
+  "n_fft": 400,
+  "n_samples": 480000,
+  "nb_max_frames": 3000,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "return_attention_mask": false,
+  "sampling_rate": 16000
+}

processing_arkasr.py

@@ -0,0 +1,445 @@
+# 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 = "<<AUDIO_TOKENS>>"
+
+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"]

processor_config.json

@@ -0,0 +1,20 @@
+{
+    "processor_class": "ArkasrProcessor",
+    "auto_map": {
+      "AutoProcessor": "processing_arkasr.ArkasrProcessor"
+    },
+  
+    "feature_extractor_type": "WhisperFeatureExtractor",
+    "tokenizer_class": "Qwen2Tokenizer",
+  
+    "merge_factor": 4,
+    "audio_token": "<|audio|>",
+  
+    "audio_dtype": "bfloat16",
+  
+    "tokenizer_config": {
+      "padding_side": "left",
+      "model_max_length": 8192
+    }
+  }
+  

spark_tokenizer_model/config.json

@@ -0,0 +1,83 @@
+{
+  "activation_dropout": 0.0,
+  "apply_spec_augment": true,
+  "architectures": [
+    "Wav2Vec2ForPreTraining"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": 1,
+  "codevector_dim": 768,
+  "contrastive_logits_temperature": 0.1,
+  "conv_bias": true,
+  "conv_dim": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "conv_kernel": [
+    10,
+    3,
+    3,
+    3,
+    3,
+    2,
+    2
+  ],
+  "conv_stride": [
+    5,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2
+  ],
+  "ctc_loss_reduction": "sum",
+  "ctc_zero_infinity": false,
+  "diversity_loss_weight": 0.1,
+  "do_stable_layer_norm": true,
+  "eos_token_id": 2,
+  "feat_extract_activation": "gelu",
+  "feat_extract_dropout": 0.0,
+  "feat_extract_norm": "layer",
+  "feat_proj_dropout": 0.1,
+  "feat_quantizer_dropout": 0.0,
+  "final_dropout": 0.0,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "layerdrop": 0.1,
+  "mask_channel_length": 10,
+  "mask_channel_min_space": 1,
+  "mask_channel_other": 0.0,
+  "mask_channel_prob": 0.0,
+  "mask_channel_selection": "static",
+  "mask_feature_length": 10,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_space": 1,
+  "mask_time_other": 0.0,
+  "mask_time_prob": 0.075,
+  "mask_time_selection": "static",
+  "model_type": "wav2vec2",
+  "num_attention_heads": 16,
+  "num_codevector_groups": 2,
+  "num_codevectors_per_group": 320,
+  "num_conv_pos_embedding_groups": 16,
+  "num_conv_pos_embeddings": 128,
+  "num_feat_extract_layers": 7,
+  "num_hidden_layers": 24,
+  "num_negatives": 100,
+  "pad_token_id": 0,
+  "proj_codevector_dim": 768,
+  "transformers_version": "4.7.0.dev0",
+  "vocab_size": 32
+}

spark_tokenizer_model/config.yaml

@@ -0,0 +1,66 @@
+audio_tokenizer:
+  mel_params:
+    sample_rate: 16000
+    n_fft: 1024
+    win_length: 640
+    hop_length: 320
+    mel_fmin: 10 
+    mel_fmax: null
+    num_mels: 128
+
+  encoder:
+    input_channels: 1024
+    vocos_dim: 384
+    vocos_intermediate_dim: 2048
+    vocos_num_layers: 12
+    out_channels: 1024
+    sample_ratios: [1,1]
+
+  decoder:
+    input_channel: 1024
+    channels: 1536
+    rates: [8, 5, 4, 2]
+    kernel_sizes: [16,11,8,4]
+
+  quantizer:
+    input_dim: 1024
+    codebook_size: 8192
+    codebook_dim: 8
+    commitment: 0.25
+    codebook_loss_weight: 2.0
+    use_l2_normlize: True
+    threshold_ema_dead_code: 0.2
+  
+  speaker_encoder:
+    input_dim: 128
+    out_dim: 1024
+    latent_dim: 128
+    token_num: 32
+    fsq_levels: [4, 4, 4, 4, 4, 4]
+    fsq_num_quantizers: 1
+
+  prenet:
+    input_channels: 1024
+    vocos_dim: 384
+    vocos_intermediate_dim: 2048
+    vocos_num_layers: 12
+    out_channels: 1024
+    condition_dim: 1024
+    sample_ratios: [1,1]
+    use_tanh_at_final: False
+
+  postnet: 
+    input_channels: 1024
+    vocos_dim: 384
+    vocos_intermediate_dim: 2048
+    vocos_num_layers: 6
+    out_channels: 1024
+    use_tanh_at_final: False
+highpass_cutoff_freq: 40
+sample_rate: 16000 
+segment_duration: 2.4 # (s)
+max_val_duration: 12 # (s)
+latent_hop_length: 320
+ref_segment_duration: 6
+volume_normalize: true
+

spark_tokenizer_model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9940cd48d4446e4340ced82d234bf5618350dd9f5db900ebe47a4fdb03867ec
+size 625518756

spark_tokenizer_model/wav2vec2-large-xlsr-53/config.json

@@ -0,0 +1,83 @@
+{
+  "activation_dropout": 0.0,
+  "apply_spec_augment": true,
+  "architectures": [
+    "Wav2Vec2ForPreTraining"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": 1,
+  "codevector_dim": 768,
+  "contrastive_logits_temperature": 0.1,
+  "conv_bias": true,
+  "conv_dim": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "conv_kernel": [
+    10,
+    3,
+    3,
+    3,
+    3,
+    2,
+    2
+  ],
+  "conv_stride": [
+    5,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2
+  ],
+  "ctc_loss_reduction": "sum",
+  "ctc_zero_infinity": false,
+  "diversity_loss_weight": 0.1,
+  "do_stable_layer_norm": true,
+  "eos_token_id": 2,
+  "feat_extract_activation": "gelu",
+  "feat_extract_dropout": 0.0,
+  "feat_extract_norm": "layer",
+  "feat_proj_dropout": 0.1,
+  "feat_quantizer_dropout": 0.0,
+  "final_dropout": 0.0,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "layerdrop": 0.1,
+  "mask_channel_length": 10,
+  "mask_channel_min_space": 1,
+  "mask_channel_other": 0.0,
+  "mask_channel_prob": 0.0,
+  "mask_channel_selection": "static",
+  "mask_feature_length": 10,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_space": 1,
+  "mask_time_other": 0.0,
+  "mask_time_prob": 0.075,
+  "mask_time_selection": "static",
+  "model_type": "wav2vec2",
+  "num_attention_heads": 16,
+  "num_codevector_groups": 2,
+  "num_codevectors_per_group": 320,
+  "num_conv_pos_embedding_groups": 16,
+  "num_conv_pos_embeddings": 128,
+  "num_feat_extract_layers": 7,
+  "num_hidden_layers": 24,
+  "num_negatives": 100,
+  "pad_token_id": 0,
+  "proj_codevector_dim": 768,
+  "transformers_version": "4.7.0.dev0",
+  "vocab_size": 32
+}

spark_tokenizer_model/wav2vec2-large-xlsr-53/preprocessor_config.json

@@ -0,0 +1,9 @@
+{
+  "do_normalize": true,
+  "feature_extractor_type": "Wav2Vec2FeatureExtractor",
+  "feature_size": 1,
+  "padding_side": "right",
+  "padding_value": 0,
+  "return_attention_mask": true,
+  "sampling_rate": 16000
+}

spark_tokenizer_model/wav2vec2-large-xlsr-53/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:314340227371a608f71adcd5f0de5933824fe77e55822aa4b24dba9c1c364dcb
+size 1269737156

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "additional_special_tokens": [
+    "<|user|>",
+    "<|begin_of_audio|>",
+    "<|end_of_audio|>",
+    "<|assistant|>",
+    "<|system|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff07bfc6cf4ed2365e9ae107e5118c89170363624f2036c85a27904d368efd87
+size 13894630

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:347b151ccd4d11234353c5ba9417ccbc5ac6e3003a9a3e1481f87080d978d782
+size 7313