Upload 56 files

added_tokens.json

@@ -7,32 +7,32 @@
   "<B_USYS>": 151666,
   "<C_A>": 151668,
   "<C_Q>": 151667,
-  "<audio_delim_baichuan>": 151693,
-  "<audio_end_baichuan>": 151677,
-  "<audio_pad_baichuan>": 151678,
-  "<audio_start_baichuan>": 151676,
-  "<baichuan_pad_token>": 151691,
-  "<box_delim_baichuan>": 151685,
-  "<box_end_baichuan>": 151684,
-  "<box_start_baichuan>": 151683,
+  "<audio_delim_ocean>": 151693,
+  "<audio_end_ocean>": 151677,
+  "<audio_pad_ocean>": 151678,
+  "<audio_start_ocean>": 151676,
+  "<ocean_pad_token>": 151691,
+  "<box_delim_ocean>": 151685,
+  "<box_end_ocean>": 151684,
+  "<box_start_ocean>": 151683,
   "<calc_end>": 151674,
   "<calc_start>": 151673,
   "<function_calling>": 151672,
-  "<img_delim_baichuan>": 151688,
-  "<img_end_baichuan>": 151680,
-  "<img_newline_baichuan>": 151682,
-  "<img_pad_baichuan>": 151681,
-  "<img_start_baichuan>": 151679,
+  "<img_delim_ocean>": 151688,
+  "<img_end_ocean>": 151680,
+  "<img_newline_ocean>": 151682,
+  "<img_pad_ocean>": 151681,
+  "<img_start_ocean>": 151679,
   "<inner_think>": 151675,
-  "<polygon_end_baichuan>": 151690,
-  "<polygon_start_baichuan>": 151689,
-  "<ref_end_baichuan>": 151687,
-  "<ref_start_baichuan>": 151686,
+  "<polygon_end_ocean>": 151690,
+  "<polygon_start_ocean>": 151689,
+  "<ref_end_ocean>": 151687,
+  "<ref_start_ocean>": 151686,
   "<reserved_113>": 151692,
   "<tool_call>": 151657,
-  "<video_end_baichuan>": 151696,
-  "<video_palce_baichuan>": 151694,
-  "<video_start_baichuan>": 151695,
+  "<video_end_ocean>": 151696,
+  "<video_palce_ocean>": 151694,
+  "<video_start_ocean>": 151695,
   "<|box_end|>": 151649,
   "<|box_start|>": 151648,
   "<|endoftext|>": 151643,

audio_modeling_ocean.py

@@ -0,0 +1,194 @@
+import torch, random, fire
+from transformers.models.whisper import WhisperConfig
+from torch.nn import functional as F
+from flash_attn import flash_attn_varlen_func
+from torch import nn
+import numpy as np
+from transformers.activations import ACT2FN
+import math
+
+def sinusoids(length, channels, max_timescale=10000):
+    """Returns sinusoids for positional embedding"""
+    assert channels % 2 == 0
+    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+
+class OceanWhisperAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor):
+        bsz, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states).view(bsz, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, self.num_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, self.num_heads, self.head_dim)
+
+        cu_len = F.pad(torch.cumsum(seq_len, dim=0), (1, 0), "constant", 0).to(torch.int32)
+        max_seqlen = torch.max(seq_len).to(torch.int32).detach()
+        attn_output = flash_attn_varlen_func(query_states, key_states, value_states, cu_len, cu_len, max_seqlen, max_seqlen, causal=False)  # (bsz * qlen, nheads, headdim)
+        attn_output = attn_output.reshape(bsz, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+        return attn_output
+
+class OceanWhisperEncoderLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embed_dim = config.d_model
+        self.self_attn = OceanWhisperAttention(self.embed_dim, config.encoder_attention_heads)
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+    def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, seq_len)
+        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 = self.fc2(hidden_states)
+        hidden_states = residual + hidden_states
+
+        if (hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16) and (
+            torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
+        ):
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+        return hidden_states
+
+class OceanAudioEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        config._attn_implementation = 'flash_attention_2'  #
+        self.config = config
+        self.max_source_positions = (config.max_audio_seconds * config.sampling_rate // config.hop_length) // config.stride_size
+        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+        # 需要在LLM的初始化中注册注册
+        self.conv1 = nn.Conv1d(config.num_mel_bins, config.d_model, kernel_size=config.kernel_size, padding=1)
+        self.conv2 = nn.Conv1d(config.d_model, config.d_model, kernel_size=config.kernel_size, stride=config.stride_size, padding=1)
+        self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, config.d_model))  # 1500 * d
+
+        self.layers = nn.ModuleList([OceanWhisperEncoderLayer(config) for _ in range(config.encoder_layers)])
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+        self.gradient_checkpointing = True
+
+    @torch.no_grad()
+    def fake_input(self, device):
+        input_features = torch.rand([2, self.config.num_mel_bins, 10], dtype=torch.float32, device=device)
+        encoder_length = torch.ones([2], dtype=torch.int32, device=device) * 3
+        bridge_length = torch.ones([2], dtype=torch.int32, device=device)
+        return input_features, encoder_length, bridge_length
+
+    def forward(
+        self,
+        input_features,
+        output_length,  # MAKESURE 输入的必须是两次conv计算后的hidden state长度
+    ):
+        input_features = input_features.to(self.conv1.weight.dtype)
+        inputs_embeds = nn.functional.gelu(self.conv1(input_features))  # (bs, channels, frames)
+        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))  # (bs, channels, frames // 2)
+        inputs_embeds = inputs_embeds.permute(0, 2, 1)  # (bs, frams, channels) 
+        bsz, tgt_len, _ = inputs_embeds.size()  # 当前batch最大长度
+        if tgt_len < self.positional_embedding.shape[0]:  
+            current_positional_embedding = self.positional_embedding[:tgt_len]
+        else:
+            current_positional_embedding = self.positional_embedding
+        hidden_states = (inputs_embeds.to(torch.float32) + current_positional_embedding).to(inputs_embeds.dtype)
+
+        # packing hidden states
+        attention_mask = torch.arange(0, tgt_len).to(hidden_states.device)
+        attention_mask = torch.lt(attention_mask, output_length.reshape(bsz, 1)).view(bsz, tgt_len, 1)
+        unpacking_index = torch.cumsum(attention_mask.to(torch.int32).view(-1), dim=0) - 1  # 转成下标
+        hidden_states = torch.masked_select(hidden_states, attention_mask).view(torch.sum(output_length), self.config.d_model)
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if self.gradient_checkpointing and self.training:
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+                    return custom_forward
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(encoder_layer),
+                    hidden_states,
+                    output_length
+                )
+            else:
+                hidden_states = encoder_layer(hidden_states, output_length)
+        hidden_states = self.layer_norm(hidden_states)
+        # unpacking
+        hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(bsz, tgt_len, self.config.d_model)
+        hidden_states = torch.where(attention_mask, hidden_states, 0)
+        return hidden_states
+
+class OceanAudioBridge(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config.audio_config
+        if self.config.avg_pooler > 1:
+            self.avg_pooler = nn.AvgPool1d(self.config.avg_pooler, stride=2)
+        else:
+            self.avg_pooler = None
+        self.proj1 = nn.Linear(self.config.d_model, config.hidden_size)
+        self.proj2 = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, x, output_length): 
+        if self.avg_pooler is not None:
+            x = x.permute(0, 2, 1)
+            x = self.avg_pooler(x)
+            x = x.permute(0, 2, 1)
+
+        batch_size, sl, _ = x.shape
+        output_length = output_length.to(x.device)
+        valid_mask = torch.arange(0, sl).to(x.device) 
+        valid_mask = torch.lt(valid_mask, output_length.reshape(batch_size, 1)).reshape(batch_size, sl, 1)
+        x = torch.masked_select(x, valid_mask).reshape(-1, self.config.d_model) # (sum(valid_sequence_length), d)
+        x = ACT2FN[self.config.activation_function](self.proj1(x))
+        x = self.proj2(x)
+        return x
+
+
+def test_audio():
+    from transformers import AutoConfig
+    from processor_ocean import OceanAudioProcessor
+    # from ..configuration_ocean import OceanConfig
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+
+    config.audio_config.d_model = 24
+    config.audio_config.encoder_layers = 2
+    config.audio_config.encoder_attention_heads = 4
+    config.audio_config.encoder_ffn_dim = 48
+
+    ae = OceanAudioEncoder(config.audio_config).cuda().to(torch.bfloat16)
+    bg = OceanAudioBridge(config).cuda().to(torch.bfloat16)
+    l = random.randint(10, 30)
+    bs = 3
+    input_length = torch.tensor([random.randint(1, l) for _ in range(bs)])
+    encoder_length, bridge_length = OceanAudioProcessor.inference_output_length(config.audio_config, input_length)
+    print("l={}, input_valid_length={},\nencoder_valid_length={}, bridge_valid_length={}".format(l, input_length, encoder_length, bridge_length))
+    wave_features = torch.rand((bs, config.audio_config.num_mel_bins, l))
+    a = ae(wave_features.to('cuda'), encoder_length.to('cuda'))
+    b = bg(a, bridge_length.to('cuda'))
+    print('encoder output={}, bridge output={}'.format(a.shape, b.shape))
+    print(a)
+    print(b)
+
+
+if __name__ == '__main__':
+    fire.Fire()
+    

config.json

@@ -1,7 +1,7 @@
 {
   "_name_or_path": "/cpfs/29f69eb5e2e60f26/code/mllm/zhangtao02/workspace/3_bc_mllm/models/mm_pretrain/cs_ocr3b/cs_ocr3b_ift_1204_ckpt_epoch_1_12030256",
   "architectures": [
-    "BaichuanForCausalLM"
+    "OceanForCausalLM"
   ],
   "attention_qkv_bias": true,
   "attention_qkv_pack": false,
@@ -117,10 +117,10 @@
     "vocab_size": 51865
   },
   "auto_map": {
-    "AutoConfig": "configuration_baichuan.BaichuanConfig",
-    "AutoModelForCausalLM": "modeling_baichuan.BaichuanForCausalLM"
+    "AutoConfig": "configuration_ocean.OceanConfig",
+    "AutoModelForCausalLM": "modeling_ocean.OceanForCausalLM"
   },
-  "baichuan_tokenizer_type": "auto",
+  "ocean_tokenizer_type": "auto",
   "bos_token_id": 1,
   "eos_token_id": 2,
   "head_dim": 128,
@@ -130,7 +130,7 @@
   "intermediate_size": 11008,
   "max_position_embeddings": 8192,
   "max_window_layers": 36,
-  "model_type": "baichuan",
+  "model_type": "ocean",
   "moe": false,
   "multimodal": [
     "image"

configuration_ocean.py

@@ -0,0 +1,111 @@
+# Copyright 2023 Ocean Inc. All Rights Reserved.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers import WhisperConfig
+from transformers import CLIPVisionConfig
+
+logger = logging.get_logger(__name__)
+
+
+class OceanConfig(PretrainedConfig):
+    model_type = "ocean"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=125696,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        sparse_attention_heads=None,
+        sparse_attention_layers=[],
+        head_dim=None,
+        attention_qkv_pack=True,
+        attention_qkv_bias=False,
+        use_norm_head=True,
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        position_embedding_type="rope",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        audio_config=None,
+        visual_config=None,
+        video_config=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads or self.num_attention_heads
+        self.sparse_attention_heads = sparse_attention_heads
+        self.sparse_attention_layers = sparse_attention_layers
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.attention_qkv_pack = attention_qkv_pack
+        self.attention_qkv_bias = attention_qkv_bias
+        self.use_norm_head = use_norm_head
+        self.hidden_act = hidden_act
+        self.position_embedding_type = position_embedding_type
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        assert self.position_embedding_type.lower() in ("rope", "alibi")
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        if audio_config is not None:
+            self.audio_config = WhisperConfig(**audio_config)
+        if visual_config is not None:
+            self.visual_config = CLIPVisionConfig(**visual_config)
+        if video_config is not None:
+            self.video_config = CLIPVisionConfig(**video_config)
+
+
+    def to_diff_dict(self):
+        data = super().to_diff_dict()
+        data["model_type"] = self.model_type
+        return data
+
+    def get_rotary_base(self):
+        if hasattr(self, "rotary_emb_base"):
+            return self.rotary_emb_base
+        else:
+            return self.rope_theta
+
+if __name__ == '__main__':  
+    from transformers import AutoConfig
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    print(config)

modeling_ocean.py

@@ -0,0 +1,1001 @@
+# Copyright 2023 Ocean Inc. All Rights Reserved.
+#
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch Ocean model."""
+import os
+import json
+import math
+from typing import List, Optional, Tuple, Union
+from threading import Thread
+from easydict import EasyDict
+import numpy as np
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from torch.nn import functional as F
+
+from transformers import PreTrainedModel
+from transformers.activations import ACT2FN
+from dataclasses import dataclass
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput
+from transformers.generation.utils import GenerationConfig
+from transformers.utils import logging
+
+from .configuration_ocean import OceanConfig
+from .audio_modeling_ocean import OceanAudioEncoder, OceanAudioBridge
+from .visual_modeling_ocean import OceanVisualEncoder, OceanVisualBridge
+from .processor_ocean import OceanMMProcessor
+from .moe import moe_matmul
+
+# support model path contain point(.)
+try:
+    # step1: copy relative imports to transformers_modules
+    from .generation_utils import build_chat_input, TextIterStreamer
+    from .sequence_parallel_utils import (
+        create_attention_layer,
+        get_sequence_parallel_size,
+        get_sequence_parallel_chunk,
+    )
+except ModuleNotFoundError:
+    # step2: direct import from transformers_modules
+    try:  # bypass check_imports failure
+        import sys
+        sys.path.append(os.path.dirname(__file__))
+        from generation_utils import build_chat_input, TextIterStreamer
+        from sequence_parallel_utils import (
+            create_attention_layer,
+            get_sequence_parallel_size,
+            get_sequence_parallel_chunk,
+        )
+    except Exception:
+        raise
+
+
+logger = logging.get_logger(__name__)
+
+def get_slopes(n):
+    def get_slopes_power_of_2(n):
+        start = (2 ** (-2 ** -(math.log2(n) - 3)))
+        ratio = start
+        return [start * ratio ** i for i in range(n)]
+
+    if math.log2(n).is_integer():
+        return get_slopes_power_of_2(
+            n)  # In the paper, we only train models that have 2^a heads for some a. This function has
+    else:  # some good properties that only occur when the input is a power of 2. To maintain that even
+        closest_power_of_2 = 2 ** math.floor(
+            math.log2(n))  # when the number of heads is not a power of 2, we use this workaround.
+        return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][
+                                                           :n - closest_power_of_2]
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        # convert into half-precision if necessary
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+
+        return self.weight * hidden_states
+
+
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=5e6, device=None):
+        super().__init__()
+        # 修复RePE初始化精度问题 https://zhuanlan.zhihu.com/p/678963442
+        # DeepSpeed 会 Hack torch.arange 强制在 GPU 上运行，这里使用原生的 torch.arange
+        try:
+            import deepspeed
+            self.arange = deepspeed.runtime.zero.partition_parameters._orig_torch_arange
+        except:
+            self.arange = torch.arange
+
+        self.inv_freq = 1.0 / (base ** (self.arange(0, dim, 2).float().to(device) / dim))
+        self.max_seq_len_cached = max_position_embeddings
+        t = self.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32)
+        self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+            t = self.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
+            freqs = torch.outer(t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(x.device)
+            self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(x.device)
+        return (
+            self.cos_cached[:, :, :seq_len, ...].to(torch.float32).to(x.device),
+            self.sin_cached[:, :, :seq_len, ...].to(torch.float32).to(x.device),
+        )
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos_, sin_, position_ids):
+    cos = cos_.squeeze(1).squeeze(0)  # [seq_len, dim]
+    sin = sin_.squeeze(1).squeeze(0)  # [seq_len, dim]
+    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+    ):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    def __init__(self, config: OceanConfig, is_sparse=False):
+        super().__init__()
+        self.config = config
+        self.position_embedding_type = config.position_embedding_type.lower()
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.hidden_size = config.num_attention_heads * self.head_dim
+        self.hidden_kv_size = self.num_kv_heads * self.head_dim
+
+        if is_sparse:
+            self.num_heads = config.sparse_attention_heads
+            assert self.num_kv_heads == config.num_attention_heads
+            self.W_pack = nn.Linear(self.hidden_size, 3 * self.num_heads * self.head_dim, bias=config.attention_qkv_bias)
+            self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        else:
+            self.num_heads = config.num_attention_heads
+            if self.config.attention_qkv_pack:
+                self.W_pack = nn.Linear(config.hidden_size, self.hidden_size + self.hidden_kv_size * 2, bias=config.attention_qkv_bias)
+                if config.moe:
+                    self.moe_W_pack = nn.Linear(config.hidden_size, self.hidden_size + self.hidden_kv_size * 2, bias=False)
+            else:
+                self.q_proj = nn.Linear(config.hidden_size, self.hidden_size, bias=config.attention_qkv_bias)
+                self.k_proj = nn.Linear(config.hidden_size, self.hidden_kv_size, bias=config.attention_qkv_bias)
+                self.v_proj = nn.Linear(config.hidden_size, self.hidden_kv_size, bias=config.attention_qkv_bias)
+
+            self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+            if config.moe:
+                self.moe_o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+
+        if self.position_embedding_type == 'rope':
+            self.rotary_emb = RotaryEmbedding(
+                dim=self.head_dim,
+                max_position_embeddings=config.max_position_embeddings,
+                base=config.get_rotary_base()
+            )
+        elif self.position_embedding_type == 'alibi':
+            self.alibi_slopes = get_slopes(self.num_heads)
+        self.attention = create_attention_layer(self.hidden_size, self.num_heads, self.head_dim)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def _repeat_kv(self, hidden_states: torch.Tensor, num_heads: int) -> torch.Tensor:
+        assert hidden_states.size(1) <= num_heads and num_heads % hidden_states.size(1) == 0
+        return repeat_kv(hidden_states, num_heads // hidden_states.size(1))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        seqlens: Optional[torch.IntTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        group_index=None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len = hidden_states.shape[:2]
+
+        if self.config.attention_qkv_pack:
+            if self.config.moe and group_index is not None:
+                proj = moe_matmul(hidden_states, [self.W_pack.weight, self.moe_W_pack.weight], group_index, lambda x, y: torch.einsum('bd,ld->bl', x, y))
+                if self.config.attention_qkv_bias:
+                    proj += self.W_pack.bias
+            else:
+                proj = self.W_pack(hidden_states)
+            query_states, key_states, value_states = proj.split([self.hidden_size, self.hidden_kv_size, self.hidden_kv_size], dim=-1)
+        else:
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+        # (B, S, hidden_size) -> (B, num_heads, S, head_size)
+        query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+        # (B, S, hidden_size) -> (B, num_kv_heads, S, head_size)
+        key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+        if self.position_embedding_type == 'rope':
+            cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len * get_sequence_parallel_size())
+            query_states, key_states = apply_rotary_pos_emb(
+                query_states, key_states, cos, sin,
+                get_sequence_parallel_chunk(position_ids)
+            )
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # repeat k/v heads if n_kv_heads < n_heads
+        key_states = self._repeat_kv(key_states, query_states.size(1))
+        value_states = self._repeat_kv(value_states, query_states.size(1))
+
+        if seqlens is not None:
+            seqlens = seqlens.to(dtype=torch.int32)
+            max_seqlen = (seqlens[1:] - seqlens[:-1]).max().item()
+            if self.position_embedding_type == 'alibi':
+                alibi_slopes = torch.tensor(self.alibi_slopes, dtype=torch.float32).to(query_states.device)
+            else:
+                alibi_slopes = None
+            attn_output = self.attention(
+                query_states, key_states, value_states, seqlens, seqlens,
+                max_seqlen, max_seqlen, causal=True, alibi_slopes=alibi_slopes, use_flash=True)
+        else:
+            attn_output = self.attention(
+                query_states, key_states, value_states, attn_mask=attention_mask, use_flash=False)
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        if not self.config.moe or group_index is None:
+            attn_output = self.o_proj(attn_output)
+        else:
+            attn_output = moe_matmul(attn_output, [self.o_proj.weight, self.moe_o_proj.weight], group_index, lambda x, y: torch.einsum('bd,ld->bl', x, y))
+
+        return attn_output, None, past_key_value
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: OceanConfig, is_sparse=False):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Attention(config=config, is_sparse=is_sparse)
+        self.mlp = MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        seqlens: Optional[torch.IntTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        group_index=None,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            seqlens=seqlens,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            group_index=group_index,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+class OceanPreTrainedModel(PreTrainedModel):
+    config_class = OceanConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DecoderLayer"]
+    _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear) or isinstance(module, nn.Conv1d):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, OceanModel):
+            module.gradient_checkpointing = value
+
+
+class OceanModel(OceanPreTrainedModel):
+    def __init__(self, config: OceanConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.merge_size = 1
+        if config.audio_config.enable:
+            self.audio_model = OceanAudioEncoder(config.audio_config)
+            self.audio_bridge_model = OceanAudioBridge(config)
+        if config.visual_config.enable:
+            self.visual_model = OceanVisualEncoder(config.visual_config)
+            self.visual_bridge_model = OceanVisualBridge(config.visual_config)
+            self.merge_size = max(config.visual_config.merge_size, self.merge_size)
+        if config.video_config.enable: # in case 没有visual_config而只有video_config
+            if not config.visual_config.enable:
+                self.visual_model = OceanVisualEncoder(config.video_config)
+            self.video_bridge_model = OceanVisualBridge(config.video_config)
+            self.merge_size = max(config.video_config.merge_size, self.merge_size)
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([
+            DecoderLayer(config, is_sparse=layer_idx in config.sparse_attention_layers)
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = True
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def get_multimodal_mask(self, input_ids, pad_token_id, special_token_list):
+        '''
+        获取任意模态的特殊mask，包含以下
+        1. pad mask 表示文本中图像/语音/视频模态提前留出的token位置
+        2. special token mask 特殊token 例如对理解模型<start> <end> 不需要next token prediction
+        3. embedding mask / lm_head mask 标记出特殊token在embedding中的mask
+        '''
+
+        pad_mask = torch.eq(input_ids, pad_token_id)
+        sp_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        lm_head_mask = torch.zeros([self.config.vocab_size, 1], dtype=torch.bool)
+        for sp_id in special_token_list:
+            sp_mask = torch.logical_or(sp_mask, torch.eq(input_ids, sp_id))
+            lm_head_mask[sp_id, 0] = True
+        return pad_mask, sp_mask, lm_head_mask
+
+    def get_audio_embed(
+            self, 
+            input_ids,
+            text_embedding,  # 1. self.embed_tokens(input_ids) 2. 其他模态结果
+            features,  # list of tensors
+            encoder_length,
+            bridge_length, 
+            group_index=None,  # 某种模态的编号 for MoE
+        ):
+        pad_mask, sp_mask, _ = self.get_multimodal_mask(input_ids, self.config.audio_config.audio_pad_token_id, self.config.multimodal_special_token_list)
+        if features is None or len(features) <= 0 : # 空list or None 保证梯度回传
+            features, encoder_length, bridge_length = self.audio_model.fake_input(input_ids.device)
+            fake_input = True
+        else:
+            fake_input = False
+        audio_embed = self.audio_model(features, encoder_length)
+        audio_embed = self.audio_bridge_model(audio_embed, bridge_length)  # (?, d)
+        if not self.training:  # 推理支持auto map 把多模态模块输出和input_ids 统一到一个device
+            audio_embed = audio_embed.to(input_ids.device)
+        if not fake_input:  # 检查多模态token 和 pad mask数量一致 （不正确的截断会导致该问题）
+            assert pad_mask.sum() == audio_embed.shape[0]
+        else:
+            assert pad_mask.sum() <= 0  # 0 vs 1
+
+        # 合并 当前模态embeddings 和text embeddings
+        input_ids = torch.where(pad_mask, torch.cumsum(pad_mask.view(-1).to(input_ids), dim=0).view(input_ids.shape)-1, input_ids)
+        if self.config.train_multimodal_special_tokens_only and self.training:  
+            # 仅special token传梯度到embedding weight, 保证LLM部分不变
+            # 注意: 多种模态之间special token list应该共享，否则会有部分被stop gradient
+            sp_mask = sp_mask.unsqueeze(-1).to(text_embedding)
+            text_embedding = (1 - sp_mask) * text_embedding.detach() + sp_mask * text_embedding 
+        text_embedding = (1 - pad_mask.to(text_embedding)).unsqueeze(-1) * text_embedding  # pad token位置填0 (不传梯度)
+        multimodal_embedding = torch.embedding(audio_embed, input_ids * pad_mask)  # 非 pad token 位置填idx=0位置结果
+        multimodal_embedding = pad_mask.to(multimodal_embedding).unsqueeze(-1) * multimodal_embedding  # 非pad token 位置填0
+
+        final_embedding = multimodal_embedding.to(text_embedding) + text_embedding
+
+        if group_index is None:
+            group_index = pad_mask.to(torch.int32)
+        else:
+            current_index = torch.max(group_index) + 1
+            group_index += pad_mask.to(torch.int32) * current_index  # 假设模态无重叠
+
+        return final_embedding, group_index  # group_index 不传None 防止MoE部分参数无梯度
+
+    def get_visual_embed(
+            self, 
+            input_ids,
+            text_embedding,  # 1. self.embed_tokens(input_ids) 2. 其他模态结果
+            images,
+            group_index,  # 某种模态的编号 for MoE
+            images_grid
+        ): 
+        # TODO 与get_audio_embed合并重复功能 减少冗余代码
+        pad_mask, sp_mask, _ = self.get_multimodal_mask(input_ids, self.config.visual_config.image_pad_token_id, self.config.multimodal_special_token_list)
+        if images is None or len(images) <= 0 : # 空list or None 保证梯度回传
+            images = self.visual_model.fake_input(input_ids.device, self.merge_size)
+            images_grid = [(1, self.merge_size, self.merge_size)]
+            fake_input = True
+        else:
+            fake_input = False
+        
+        images = torch.cat(images, dim=0)
+        images_grid = torch.tensor(np.array(images_grid))
+        visual_embed = self.visual_model(images, grid_thw=images_grid)
+        visual_embed = self.visual_bridge_model(visual_embed)
+
+        if not self.training:  # 推理支持auto map 把多模态模块输出和input_ids 统一到一个device
+            visual_embed = visual_embed.to(input_ids.device)
+        if not fake_input:  # 检查多模态token 和 pad mask数量一致 （不正确的截断会导致该问题）
+            assert pad_mask.sum() == visual_embed.shape[0], '{} != {}'.format(pad_mask.sum(), visual_embed.shape[0])
+        else:
+            assert pad_mask.sum() <= 0, '{} != {}'.format(pad_mask.sum(), visual_embed.shape[0])
+
+        # 合并 当前模态embeddings 和text embeddings
+        input_ids = torch.where(pad_mask, torch.cumsum(pad_mask.view(-1).to(input_ids), dim=0).view(input_ids.shape)-1, input_ids)
+        if self.config.train_multimodal_special_tokens_only and self.training:  
+            # 仅special token传梯度到embedding weight, 保证LLM部分不变
+            # 注意: 多种模态之间special token list应���共享，否则会有部分被stop gradient
+            sp_mask = sp_mask.unsqueeze(-1).to(text_embedding)
+            text_embedding = (1 - sp_mask) * text_embedding.detach() + sp_mask * text_embedding 
+        text_embedding = (1 - pad_mask.to(text_embedding)).unsqueeze(-1) * text_embedding  # pad token位置填0 (不传梯度)
+        multimodal_embedding = torch.embedding(visual_embed, input_ids * pad_mask)  # 非 pad token 位置填idx=0位置结果
+        multimodal_embedding = pad_mask.to(multimodal_embedding).unsqueeze(-1) * multimodal_embedding  # 非pad token 位置填0
+
+        final_embedding = multimodal_embedding.to(text_embedding) + text_embedding
+
+        if group_index is None:
+            group_index = pad_mask.to(torch.int32)
+        else:
+            current_index = torch.max(group_index) + 1
+            group_index += pad_mask.to(torch.int32) * current_index  # 假设模态无重叠
+
+        return final_embedding, group_index  # group_index 不传None 防止MoE部分参数无梯度
+
+    def get_video_embed(
+            self, 
+            input_ids,
+            text_embedding,  # 1. self.embed_tokens(input_ids) 2. 其他模态结果
+            images,
+            group_index,  # 某种模态的编号 for MoE
+            images_grid
+
+        ): 
+        # TODO 与get_audio_embed合并重复功能 减少冗余代码
+        pad_mask, sp_mask, _ = self.get_multimodal_mask(input_ids, self.config.video_config.video_place_token_id, self.config.multimodal_special_token_list)
+        if images is None or len(images) <= 0 : # 空list or None 保证梯度回传
+            images = self.visual_model.fake_input(input_ids.device, self.merge_size)
+            images_grid = [(1, self.merge_size, self.merge_size)]
+            fake_input = True
+        else:
+            fake_input = False
+        
+        images = torch.cat(images, dim=0)
+        images_grid = torch.tensor(np.array(images_grid))
+        visual_embed = self.visual_model(images, grid_thw=images_grid)
+        visual_embed = self.video_bridge_model(visual_embed)
+        
+        if not self.training:  # 推理支持auto map 把多模态模块输出和input_ids 统一到一个device
+            visual_embed = visual_embed.to(input_ids.device)
+        if not fake_input:  # 检查多模态token 和 pad mask数量一致 （不正确的截断会导致该问题）
+            assert pad_mask.sum() == visual_embed.shape[0], '{} != {}'.format(pad_mask.sum(), visual_embed.shape[0])
+            assert pad_mask.sum() == visual_embed.shape[0], '{} != {}'.format(pad_mask.sum(), visual_embed.shape[0])
+        else:
+            assert pad_mask.sum() <= 0, '{} != {}'.format(pad_mask.sum(), visual_embed.shape[0])
+
+        # 合并 当前模态embeddings 和text embeddings
+        input_ids = torch.where(pad_mask, torch.cumsum(pad_mask.view(-1).to(input_ids), dim=0).view(input_ids.shape)-1, input_ids)
+        if self.config.train_multimodal_special_tokens_only and self.training:  
+            # 仅special token传梯度到embedding weight, 保证LLM部分不变
+            # 注意: 多种模态之间special token list应该共享，否则会有部分被stop gradient
+            sp_mask = sp_mask.unsqueeze(-1).to(text_embedding)
+            text_embedding = (1 - sp_mask) * text_embedding.detach() + sp_mask * text_embedding 
+        text_embedding = (1 - pad_mask.to(text_embedding)).unsqueeze(-1) * text_embedding  # pad token位置填0 (不传梯度)
+        multimodal_embedding = torch.embedding(visual_embed, input_ids * pad_mask)  # 非 pad token 位置填idx=0位置结果
+        multimodal_embedding = pad_mask.to(multimodal_embedding).unsqueeze(-1) * multimodal_embedding  # 非pad token 位置填0
+
+        final_embedding = multimodal_embedding.to(text_embedding) + text_embedding
+
+        if group_index is None:
+            group_index = pad_mask.to(torch.int32)
+        else:
+            current_index = torch.max(group_index) + 1
+            group_index += pad_mask.to(torch.int32) * current_index  # 假设模态无重叠
+
+        return final_embedding, group_index  # group_index 不传None 防止MoE部分参数无梯度
+
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        seqlens: Optional[torch.IntTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        audios: Optional[List|torch.Tensor] = None,
+        encoder_length: Optional[torch.Tensor] = None,
+        bridge_length: Optional[torch.Tensor] = None,
+        images: Optional[List|torch.Tensor] = None,
+        images_grid: Optional[List|torch.Tensor] = None,
+        videos: Optional[List|torch.Tensor] = None,
+        videos_grid: Optional[List|torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        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
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        group_index = None
+        if inputs_embeds is None:
+            sp_input_ids = get_sequence_parallel_chunk(input_ids)
+            inputs_embeds = self.embed_tokens(sp_input_ids)
+            if self.config.audio_config.enable:
+                inputs_embeds, group_index = self.get_audio_embed(sp_input_ids, inputs_embeds, audios, encoder_length, bridge_length)
+            if self.config.visual_config.enable:
+                inputs_embeds, group_index = self.get_visual_embed(sp_input_ids, inputs_embeds, images, group_index, images_grid)  # 注意更新group index
+            if self.config.video_config.enable:
+                inputs_embeds, group_index = self.get_video_embed(sp_input_ids, inputs_embeds, videos, group_index, videos_grid)  # 注意更新group index
+
+        if seqlens is not None and seqlens.ndim == 2:
+            # batch multi-pack 样本拉平
+            cu_seqlens = []
+            offset, seqlen = 0, seqlens.size(1)
+            for lens in seqlens:
+                cu_seqlens.append(offset)
+                cu_seqlens.extend((lens[(lens > 0) & (lens < seqlen)] + offset).tolist())
+                offset += seqlen
+            cu_seqlens.append(offset)
+            seqlens = torch.tensor(cu_seqlens, dtype=seqlens.dtype, device=seqlens.device)
+        elif seqlens is None and self.training:
+            # 兼容预训练场景, 此时 seqlens=None, 默认 maxlength
+            seqlens = torch.arange(
+                end=input_ids.size(0) + 1,
+                dtype=torch.int32,
+                device=input_ids.device
+            ) * input_ids.size(1)
+        if seqlens is not None:
+            attention_mask = None  # unset attention_mask to save memory
+
+        if seqlens is None and attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+            )
+        if attention_mask is not None:
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, False, group_index)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    seqlens,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    seqlens=seqlens,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    group_index=group_index,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class NormHead(nn.Module):
+    def __init__(self, hidden_size, vocab_size, bias=False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.vocab_size = vocab_size
+        self.weight = nn.Parameter(torch.empty((self.vocab_size, self.hidden_size)))
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def forward(self, hidden_states, mask=None):
+        norm_weight = nn.functional.normalize(self.weight)
+        if mask is not None:
+            mask = mask.to(norm_weight)
+            norm_weight = norm_weight * mask + (1 - mask) * norm_weight.detach()
+        return nn.functional.linear(hidden_states, norm_weight)
+
+
+    def extra_repr(self) -> str:
+        return f'in_features={self.hidden_size}, out_features={self.vocab_size}'
+
+@dataclass
+class OceanMMCausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    text_nt_loss: Optional[torch.FloatTensor] = None
+    flatten_loss: Optional[torch.FloatTensor] = None
+
+class OceanForCausalLM(OceanPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = OceanModel(config)
+        if config.use_norm_head:
+            self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False)
+        else:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def bind_processor(self, tokenizer, **kwargs):
+        self.processor = OceanMMProcessor(
+                tokenizer=tokenizer,
+                config=self.config,
+                **kwargs,
+                )
+        return self.processor
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        seqlens: Optional[torch.IntTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        audios: Optional[List|torch.Tensor] = None,
+        encoder_length: Optional[torch.Tensor] = None,
+        bridge_length: Optional[torch.Tensor] = None,
+        images: Optional[torch.Tensor] = None, 
+        images_grid: Optional[torch.Tensor] = None, 
+        videos: Optional[torch.Tensor] = None, 
+        videos_grid: Optional[torch.Tensor] = None, 
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        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
+
+        _, sp_mask, _ = self.model.get_multimodal_mask(input_ids, self.config.audio_config.audio_pad_token_id, self.config.multimodal_special_token_list)
+        # TODO 放开部分可学习的special token lmhead参数
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            seqlens=seqlens,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            audios=audios,
+            encoder_length=encoder_length,
+            bridge_length=bridge_length,
+            images=images,
+            images_grid=images_grid,
+            videos=videos,
+            videos_grid=videos_grid,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        
+        # 部分可学习的special token放开lm head梯度
+        special_with_loss_list = list(set(self.config.multimodal_special_token_list) - set(self.config.multimodal_special_token_no_loss_list))
+        _, sp_with_loss_mask, lm_head_mask = self.model.get_multimodal_mask(input_ids, self.config.audio_config.audio_pad_token_id, special_with_loss_list)
+        if self.config.train_multimodal_special_tokens_only and self.training and len(special_with_loss_list) > 0:
+            if self.config.use_norm_head:
+                logits = self.lm_head(hidden_states, mask=lm_head_mask)
+            else:
+                lm_head_mask = lm_head_mask.to(self.lm_head.weight)
+                norm_weight = self.lm_head.weight * lm_head_mask + (1 - lm_head_mask) * self.lm_head.weight.detach()
+                logits = torch.einsum('bsd,ld->bsl', hidden_states, norm_weight)
+        else:
+            logits = self.lm_head(hidden_states)
+
+        loss = torch.tensor(0, device=hidden_states.device, dtype=hidden_states.dtype)
+        text_nt_loss = torch.tensor(0, device=hidden_states.device, dtype=hidden_states.dtype)
+        flatten_loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            valid_mask = torch.gt(shift_labels, -1)  # label < 0 视为pad位置
+            sp_mask = sp_mask[..., 1:].contiguous()
+            text_mask = torch.logical_and(valid_mask, torch.logical_not(sp_mask))
+            
+            # Flatten the tokens
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            flatten_loss = F.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction='none')
+
+            loss = torch.mean(torch.masked_select(flatten_loss, valid_mask.view(-1)))
+            text_nt_loss = torch.mean(torch.masked_select(flatten_loss, text_mask.view(-1))).detach()
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return OceanMMCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            text_nt_loss=text_nt_loss,
+            flatten_loss=flatten_loss
+        )
+
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        elif past_key_values is not None:
+            model_inputs = {"input_ids": input_ids}
+        else:
+            model_inputs = {"input_ids": input_ids, 
+                            "audios": kwargs.get("audios", None), "encoder_length": kwargs.get("encoder_length", None), "bridge_length": kwargs.get("bridge_length", None),
+                            "images": kwargs.get("images", None),
+                            "videos": kwargs.get("videos", None)
+                            }
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images_grid":  kwargs.get("images_grid"),
+                "videos_grid":  kwargs.get("videos_grid"),
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past
+
+    def chat(self, tokenizer, messages: List[dict], stream=False,
+             generation_config: Optional[GenerationConfig]=None):
+        generation_config = generation_config or self.generation_config
+        input_ids = build_chat_input(self, tokenizer, messages, generation_config.max_new_tokens)
+        if stream:
+            streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+            Thread(target=self.generate, kwargs=dict(
+                inputs=input_ids, streamer=streamer,
+                generation_config=generation_config,
+            )).start()
+            return streamer
+        else:
+            outputs = self.generate(input_ids, generation_config=generation_config)
+            response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
+            return response

processor_ocean.py

@@ -0,0 +1,1154 @@
+import requests
+import re, ujson, os, sys, fire, glob, random, time, json
+import numpy as np
+import io
+import torch
+from torch.utils.data import default_collate
+import torchaudio
+from typing import *
+from dataclasses import dataclass, field
+import transformers
+from transformers.modeling_outputs import ModelOutput
+from transformers.audio_utils import mel_filter_bank, spectrogram, window_function
+from functools import lru_cache
+from io import BytesIO
+from PIL import Image
+from qcloud_cos import CosConfig
+from qcloud_cos import CosS3Client
+import tos
+import concurrent.futures as cf
+from transformers.image_transforms import resize, center_crop, get_resize_output_image_size
+from transformers.image_utils import PILImageResampling
+from PIL import Image, ImageOps
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+import base64
+from decord import VideoReader, cpu
+import cv2
+import av
+import imagesize
+import math
+
+
+def smart_resize(
+    height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
+):
+    """Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+
+    """
+    # if height < factor or width < factor:
+        # raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = round(height / factor) * factor if height > factor else factor
+    w_bar = round(width / factor) * factor if width > factor else factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = math.floor(height / beta / factor) * factor
+        w_bar = math.floor(width / beta / factor) * factor
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+    return h_bar, w_bar
+
+
+def select_best_resolution(image_size, candidate_resolutions):
+    '''找到最佳的resolution 对于原图进行放缩
+        image_size 通常为ori_size e.g. (8*336, 16*336)
+        candidate_resolutions 为备选分辨率 e.g. (1*336, 4*336)
+    '''
+    try:
+        original_width, original_height = image_size
+    except:
+        pass
+    best_fit = None
+    max_effective_resolution = 0
+    min_wasted_resolution = float("inf")
+    # 从candidate_resolutions 中遍历宽和高
+    for width, height in candidate_resolutions:
+        # width / original_width 和 height / original_height 中最小的那个作为scale
+        scale = min(width / original_width, height / original_height) # e.g. scale =min (1/8, 1/4) = 1/8
+        # 放缩 original_width 和 original_height
+        downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale) # e.g. 1*336, 2*336
+        # effective_resolution 为 放缩之后的分辨率 s^2 * w * h 
+        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height) # e.g. min(1*336 * 2*336, 8*336 * 16*336)
+        # wasted_resolution 为 放缩前后分辨率的差值
+        wasted_resolution = (width * height) - effective_resolution
+        # 若 （1） 放缩之后的分辨率 比当前的max_effective_resolution更大;
+            # (2) 
+        if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
+            max_effective_resolution = effective_resolution # 更新max_effective_resolution
+            min_wasted_resolution = wasted_resolution # min_wasted_resolution 
+            best_fit = (width, height)
+
+    return best_fit
+
+
+def read_video(image_path, max_frame_number, decode_way):
+    if decode_way=='1fps':
+        try:
+            vr = VideoReader(image_path, ctx=cpu(0))
+            total_frame_num = len(vr)
+            fps = round(vr.get_avg_fps())
+            frame_idx = [i for i in range(0, len(vr), fps)]
+            frames = vr.get_batch(frame_idx).asnumpy()
+            frames = [i for i in frames]
+            cnt = len(frames)
+        except Exception as e:
+            print(image_path)
+            print('error is', e)
+            return None
+    elif decode_way=='key':
+        try: 
+            with av.open(image_path) as container:                         
+                stream = container.streams.video[0]
+                stream.codec_context.skip_frame = 'NONKEY'
+                frames = []
+                fps = int(stream.average_rate)
+                cnt = 0
+                for frame in container.decode(stream): # 关键帧存成image patch
+                    image = frame.to_image()
+                    frames.append(image)
+                    cnt += 1
+        except Exception as e:
+            print('error is', e)
+            return None
+    if frames is None or len(frames)==0:
+        return None
+    if len(frames)>max_frame_number and max_frame_number>0:
+        # 生成均匀间隔的索引
+        indices = np.linspace(0, len(frames) - 1, max_frame_number, dtype=int)
+        # 根据索引获取对应元素
+        sampled_elements = [frames[idx] for idx in indices]
+        frames = sampled_elements
+    return frames
+
+class OceanImageProcessor:
+    def __init__(self, config, **kwargs):
+        self.config = config  # visual_config
+        self.min_pixels = self.config.min_pixels if hasattr(self.config, 'min_pixels') else 56 * 56
+        self.max_pixels = self.config.max_pixels if hasattr(self.config, 'max_pixels') else 28 * 28 * 1280
+        self.patch_size = self.config.patch_size if hasattr(self.config, 'patch_size') else 14
+        self.temporal_patch_size = self.config.temporal_patch_size if hasattr(self.config, 'temporal_patch_size') else 2
+        self.merge_size = self.config.merge_size if hasattr(self.config, 'merge_size') else 2
+        self.spatial_merge_size = self.config.spatial_merge_size if hasattr(self.config, 'spatial_merge_size') else 2
+
+    def image_transform(self, strseq, return_mm_data = True):
+        image = None
+        if isinstance(strseq, str):
+            if return_mm_data:
+                image = Image.open(strseq).convert("RGB") 
+        else:
+            image = Image.open(BytesIO(strseq)).convert("RGB")
+            
+        image = np.array(image.convert("RGB")) # 这一步首先将图像转换为 RGB 格式，确保图像有三个通道（R、G、B）。然后使用 np.array() 将其转换为 NumPy 数组，方便后续处理。
+        image_org_size = image.shape[:2] # 这里保存了图像的原始大小（高度和宽度），image.shape 返回图像的形状 (高度, 宽度, 通道数)，而 image.shape[:2] 提取了前两个值，即原始的高度和宽度。这个信息可以用于后续的对比或其他处理。
+        
+        # resize, crop, scale, normalize
+        # 接受目标尺寸作为输入参数，通常是目标尺寸的短边或长边长度。例如，如果指定目标短边为 336 像素，函数会自动计算出对应的长边大小，以保持图像的宽高比。
+        # 输出一个新的尺寸，这个尺寸通常是 (宽度, 高度) 格式，用于后续的图像调整操作，如缩放或裁剪。
+        resized_height, resized_width = smart_resize(
+            image_org_size[0], image_org_size[1],
+            factor=self.patch_size * self.spatial_merge_size,
+            min_pixels=self.min_pixels,
+            max_pixels=self.max_pixels,
+        )
+        output_size = (resized_height, resized_width)
+        
+        # output_size = get_resize_output_image_size(image, self.config.crop_size, False)  # 短边resize到336
+        # 使用 resize 函数将图像调整到 output_size 大小。PILImageResampling.BICUBIC 指定使用双三次插值法来进行图像缩放，这种方法通常能够提供较好的图像质量。
+        # image: 输入的图像数据，可以是 NumPy 数组或 PIL 图像对象；output_size: 目标大小，通常是一个二元组 (宽度, 高度)。这个尺寸可以是图像的绝对大小，也可以是相对于原始图像的比例；
+        # resample: 可选的重采样方法，通常用于确定如何插值像素。例如，PILImageResampling.BICUBIC 表示使用双三次插值法，这是一种平滑的插值方法，常用于图像缩放。
+        image = resize(image, output_size, PILImageResampling.BICUBIC)
+        # 从图像中心裁剪出一个指定大小的区域，这里是一个正方形区域 self.config.crop_size x self.config.crop_size。center_crop 函数的参数 return_numpy=True 表示返回一个 NumPy 数组形式的裁剪图像。
+        # image = center_crop(image, (self.config.crop_size, self.config.crop_size), return_numpy=True)
+        img = image.transpose(2, 0, 1)
+        # 对图像进行归一化和标准化处理
+        image = (img / 255.0 - np.array(self.config.image_mean)[:, np.newaxis, np.newaxis]) / np.array(self.config.image_std)[:,np.newaxis,np.newaxis]
+        # 处理成patch
+        patches = image[np.newaxis, :]
+        if patches.shape[0] == 1:
+            patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))
+        channel = patches.shape[1]
+        grid_t = patches.shape[0] // self.temporal_patch_size
+        grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
+        patches = patches.reshape(
+            grid_t,
+            self.temporal_patch_size,
+            channel,
+            grid_h // self.spatial_merge_size,
+            self.spatial_merge_size,
+            self.patch_size,
+            grid_w // self.spatial_merge_size,
+            self.spatial_merge_size,
+            self.patch_size,
+        )
+        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+        flatten_patches = patches.reshape(
+            grid_t * grid_h * grid_w, channel * self.temporal_patch_size * self.patch_size * self.patch_size
+        )
+
+        return flatten_patches, image_org_size, (grid_t, grid_h, grid_w)
+
+
+class OceanAudioProcessor:
+    # 包含基本的音频特征抽取模块 + 输入数据解析模块 + cos请求/缓存模块
+    def __init__(
+        self,
+        config,  # audio processor config
+        **kwargs
+    ):
+        # make sure you have install 'conda install -c conda-forge 'ffmpeg<7'' for torchaudio
+        assert(len(torchaudio.list_audio_backends()) > 0)
+        self.config = config
+        self.mel_filters = mel_filter_bank(
+            num_frequency_bins=1 + self.config.n_fft // 2,
+            num_mel_filters=self.config.num_mel_bins,
+            min_frequency=0.0,
+            max_frequency=self.config.sampling_rate / 2.0,
+            sampling_rate=self.config.sampling_rate,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+
+    @staticmethod
+    def zero_mean_unit_var_norm(x):
+        return (x - x.mean()) / torch.sqrt(x.var() + 1e-8)
+
+    def load_audio_waveform(self, uri, return_tensors=True, do_normalize=False):
+        metadata = torchaudio.info(uri)  # sample_rate, num_frames, num_channels, bits_per_sample, encoding=PCM_S
+        assert(metadata.num_channels <= 2), "acoustic file with {} channels.".format(metadata.num_channels)  # whisper only accept mono channel audio
+        waveform_tensor, _ = torchaudio.load(uri, normalize=True)
+        if self.config.sampling_rate != metadata.sample_rate:
+            waveform_tensor = torchaudio.functional.resample(waveform_tensor, metadata.sample_rate, self.config.sampling_rate)
+
+        # downmix to mono channel https://trac.ffmpeg.org/wiki/AudioChannelManipulation
+        if metadata.num_channels > 1:
+            waveform_tensor = torch.mean(waveform_tensor, dim=0, keepdim=True)
+
+        # normalized to zero mean (Qwen Audio没有处理 但Whisper官方实现)
+        if do_normalize:
+            waveform_tensor = self.zero_mean_unit_var_norm(waveform_tensor)
+
+        if return_tensors:  # (channels, samples)
+            return waveform_tensor
+        else:
+            return waveform_tensor.numpy()
+    
+    def split_with_overlap(self, waveform):  # 如果长度超过最大长度限制 分割为带overlap的多段
+        channels, wave_samples = waveform.shape
+        max_audio_samples = self.config.max_audio_seconds * self.config.sampling_rate
+        if wave_samples <= max_audio_samples or self.config.split_overlap < 0:
+            return [waveform]  # 没有超出最大长度or截断逻辑 统一返回list
+        
+        split_waveform, start = [], 0
+        while start < wave_samples:  # 20240724修改 统一按秒数对齐overlap 保证不同sampling rate/n_fft/hop length配置下采到的数据是一致的
+            if start > int(self.config.sampling_rate * self.config.split_overlap):
+                start -= int(self.config.sampling_rate * self.config.split_overlap)  # 0表示没有overlap，>0 overlap对应秒数
+            end = min(start + max_audio_samples, wave_samples)
+            split_waveform.append(waveform[:, start:end])  # 注意这里可能会切割出特别短的片段 需要在预处理判断并丢弃
+            start = end
+        return split_waveform
+
+    @classmethod        
+    def inference_output_length(cls, config, input_length):
+        # for whisper + bridge
+        kernel_size = config.kernel_size
+        stride_size = config.stride_size
+        avg_pooler = config.avg_pooler
+        encoder_length = (input_length + 2 * (kernel_size // 2) - kernel_size) // 1 + 1  # conv layer1 with pad=1
+        encoder_length = (encoder_length + 2 * (kernel_size // 2) - kernel_size) // stride_size + 1  # conv layer2 with pad=1
+        if avg_pooler > 1:
+            bridge_length = encoder_length // avg_pooler
+        return encoder_length, bridge_length
+
+    def extract_fbank_features(self, waveform):
+        # ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/feature_extraction_whisper.py
+        channels, wave_samples = waveform.shape
+        assert(wave_samples >= self.config.n_fft)
+        valid_frame_nums = min(self.config.max_audio_seconds * self.config.sampling_rate // self.config.hop_length, wave_samples // self.config.hop_length + 1)
+        if wave_samples < self.config.max_audio_seconds * self.config.sampling_rate:
+            waveform = torch.nn.functional.pad(waveform, (0, self.config.max_audio_seconds * self.config.sampling_rate - wave_samples), "constant", 0)
+        else:
+            waveform = waveform[:, :self.config.max_audio_seconds * self.config.sampling_rate]
+
+        window = torch.hann_window(self.config.n_fft)
+        stft = torch.stft(waveform, self.config.n_fft, self.config.hop_length, window=window, return_complex=True)  # fft, len(wave) // n_fft // 2 + 1
+        magnitudes = stft[..., :-1].abs() ** 2
+
+        mel_filters = torch.from_numpy(self.mel_filters).type(torch.float32)
+        mel_spec = mel_filters.T @ magnitudes
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+
+        if waveform.dim() == 2:
+            max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
+            log_spec = torch.maximum(log_spec, max_val - 8.0)
+        else:
+            log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+        log_spec = (log_spec + 4.0) / 4.0
+
+        log_spec = log_spec[0].numpy()  # (channel, filters, samples) -> (filters, samples)
+        log_spec[:, valid_frame_nums:] = 0.0  # pad0 在collect时取batch内最大长度
+
+        return log_spec, valid_frame_nums
+
+    def data_augment(self, feature: np.array, input_length, training=True):
+        # reference https://arxiv.org/pdf/1904.08779
+        # run only on cpu
+        def mask_start_indices(input_length, mask_length, min_masks, mask_prob):
+            # 计算总共需要mask的span数 之后随机筛选span开始下标
+            num_masked_span = int(mask_prob * input_length / mask_length + random.random())
+            num_masked_span = max(num_masked_span, min_masks)
+            start_indices = list(range(input_length - mask_length))
+            random.shuffle(start_indices)
+            start_indices = start_indices[:num_masked_span]
+            return start_indices
+
+        if not training or (self.config.mask_time_prob <= 0 and self.config.mask_feature_prob <= 0):
+            return feature
+        if input_length < self.config.mask_time_length * self.config.mask_time_min_masks + 1:
+            return feature
+        if self.config.num_mel_bins < self.config.mask_feature_length * self.config.mask_feature_min_masks + 1: 
+            return feature
+        
+        if self.config.mask_time_prob > 0:
+            start_indices = mask_start_indices(input_length, self.config.mask_time_length, self.config.mask_time_min_masks, self.config.mask_time_prob) 
+            for start_idx in start_indices:
+                feature[:, start_idx: start_idx + self.config.mask_time_length] = 0.0
+        if self.config.mask_feature_prob > 0:
+            start_indices = mask_start_indices(self.config.num_mel_bins, self.config.mask_feature_length, self.config.mask_feature_min_masks, self.config.mask_feature_prob) 
+            for start_idx in start_indices:
+                feature[start_idx: start_idx + self.config.mask_feature_length, :] = 0.0
+
+        return feature
+
+class CosClient():
+    def __init__(self, bucket_name='crawl-pic-1317568651',
+            max_retries=2):
+        self.config = CosConfig(
+            Endpoint="cos.ap-guangzhou.myqcloud.com", 
+            # Region='ap-guangzhou', 
+            SecretId='AKIDnRpxoOghgVs0tkU3Mfv20jAMI0SRDj02',
+            SecretKey='td9tRlqiPvEJ8i27wXwBIDiy5ye6JGyS',
+            Token=None, Scheme='https', Timeout=300)
+        self.client = CosS3Client(self.config)
+        self.max_retries = max_retries
+        self.bucket_name = bucket_name
+    
+    def __call__(self, relative_path, bucket_name=None):
+        if bucket_name is None or len(bucket_name) <= 0:
+            bucket_name = self.bucket_name
+        multimodal_bytes = None 
+        for _ in range(self.max_retries):
+            try:
+                response = self.client.get_object(Bucket=bucket_name, Key=relative_path)
+                fp = response['Body'].get_raw_stream()
+                multimodal_bytes = fp.read()
+                break
+            except Exception as e:
+                time.sleep(0.01)
+                continue
+        return multimodal_bytes
+
+class TosClient(object):
+    def __init__(self):
+        ak = "AKLTYTM3MWY5MTFhNDgyNDk4YjhmYTE0ZTE3YTk5ZmU1MjU"
+        sk = "TVRRM1pUZGtaVEJqWTJJd05HSTNPR0ppWVdKa1lqYzVORFUwTlRobU1UVQ=="
+        endpoint = "tos-cn-beijing.ivolces.com"  # "tos-cn-beijing.ivolces.com"
+        region = "cn-beijing"
+        self.bucket_name = "audio-dataset"
+        self.client = tos.TosClientV2(ak, sk, endpoint, region)
+    
+    def __call__(self, path, bucket_name=None):
+        if bucket_name is None:
+            bucket_name = self.bucket_name
+        for _ in range(2):
+            try:
+                object_stream = self.client.get_object(bucket_name, path)
+                return object_stream.read()
+            except Exception as e:
+                time.sleep(0.01)
+                continue
+        return None
+
+@dataclass
+class OceanProcessorOutput(ModelOutput):  
+    input_ids: Optional["List|torch.Tensor"] = None
+    labels: Optional["List|torch.Tensor"] = None
+    attention_mask: Optional["List|torch.Tensor"] = None
+    position_ids: Optional["List|torch.Tensor"] = None
+    seqlens: Optional["List|torch.Tensor"] = None  # 需要配合Ocean Modeling使用
+    # audio fields
+    audios: Optional["List|torch.Tensor"] = None
+    encoder_length: Optional["List|torch.Tensor"] = None
+    bridge_length: Optional["List|torch.Tensor"] = None
+    # image fields
+    images: Optional["List|torch.Tensor"] = None
+    patch_nums: Optional["List|torch.Tensor"] = None
+    images_size: Optional["List|torch.Tensor"] = None
+    crop_size: Optional["List|torch.Tensor"] = None
+    images_grid: Optional["List|torch.Tensor"] = None
+    # video fields
+    videos: Optional["List|torch.Tensor"] = None
+    videos_patch_nums: Optional["List|torch.Tensor"] = None
+    videos_size: Optional["List|torch.Tensor"] = None
+    videos_crop_size: Optional["List|torch.Tensor"] = None
+    videos_grid: Optional["List|torch.Tensor"] = None
+    # processor fields
+    raw_text: Optional[str] = None
+    index: Optional[int] = None
+
+    def concatenate(self, other):  # 仅限list使用
+        def concat_one(a, b):
+            if a is None and b is None:
+                return None
+            elif a is None and b is not None:
+                return b 
+            elif a is not None and b is None: 
+                return a 
+            else: 
+                return a + b
+        return OceanProcessorOutput(
+            input_ids=concat_one(self.input_ids, other.input_ids),
+            labels=concat_one(self.labels, other.labels),
+            audios=concat_one(self.audios, other.audios),
+            encoder_length=concat_one(self.encoder_length, other.encoder_length),
+            bridge_length=concat_one(self.bridge_length, other.bridge_length), 
+            images=concat_one(self.images, other.images),
+            images_grid=concat_one(self.images_grid, other.images_grid),
+            patch_nums=concat_one(self.patch_nums, other.patch_nums),
+
+            videos=concat_one(self.videos, other.videos),
+            videos_grid=concat_one(self.videos_grid, other.videos_grid),
+            videos_patch_nums=concat_one(self.videos_patch_nums, other.videos_patch_nums),
+
+            position_ids=concat_one(self.position_ids, other.position_ids),
+            seqlens=concat_one(self.seqlens, other.seqlens),
+            images_size=concat_one(self.images_size, other.images_size)
+        )
+
+class OceanMMProcessor(object):
+    def __init__(self,
+                tokenizer: transformers.PreTrainedTokenizer,
+                config,
+                training,
+                relative_path=None,
+                **kwargs, 
+    ):
+        self.tokenizer = tokenizer
+        self.config = config
+        self.audio_processor = None
+        if hasattr(config, "audio_config"):
+            self.audio_processor = OceanAudioProcessor(config.audio_config)
+        self.visual_processor = None
+        if hasattr(config, "visual_config"):
+            self.visual_processor = OceanImageProcessor(config.visual_config)
+        self.video_processor = None
+        if hasattr(config, "video_config"):
+            self.video_processor = OceanImageProcessor(config.video_config)
+        self.training = training
+        self.relative_path = relative_path
+        self.cos_client = CosClient()
+        self.tos_client = TosClient()
+        # audio tag
+        self.audio_start_tag = None
+        self.audio_end_tag = None
+        self.audio_pad_tag = None
+        self.audio_delim_tag = None
+        if hasattr(self.config, "audio_config"):
+            self.audio_start_tag = self.tokenizer.convert_ids_to_tokens(self.config.audio_config.audio_start_token_id)
+            self.audio_end_tag = self.tokenizer.convert_ids_to_tokens(self.config.audio_config.audio_end_token_id)
+            self.audio_pad_tag = self.tokenizer.convert_ids_to_tokens(self.config.audio_config.audio_pad_token_id)
+            self.audio_delim_tag = self.tokenizer.convert_ids_to_tokens(self.config.audio_config.audio_delim_token_id)
+        # image tag
+        self.image_start_tag = None
+        self.image_end_tag = None
+        self.image_pad_tag = None
+        self.video_start_tag = None
+        self.video_end_tag = None
+        if hasattr(self.config, "visual_config"):
+            # special token for start_tag
+            self.image_start_tag = self.tokenizer.convert_ids_to_tokens(self.config.visual_config.image_start_token_id)
+            # special token for end_tag
+            self.image_end_tag = self.tokenizer.convert_ids_to_tokens(self.config.visual_config.image_end_token_id)
+            # special token for pad_tag
+            self.image_pad_tag = self.tokenizer.convert_ids_to_tokens(self.config.visual_config.image_pad_token_id)
+            self.image_line_tag = self.tokenizer.convert_ids_to_tokens(self.config.visual_config.image_line_token_id)
+            self.image_delimiter_tag = self.tokenizer.convert_ids_to_tokens(self.config.visual_config.image_delimiter_token_id) 
+        if hasattr(self.config, "video_config"):
+            self.video_start_tag = self.tokenizer.convert_ids_to_tokens(self.config.video_config.video_start_token_id)
+            self.video_end_tag = self.tokenizer.convert_ids_to_tokens(self.config.video_config.video_end_token_id)
+            self.image_start_tag = self.tokenizer.convert_ids_to_tokens(self.config.video_config.image_start_token_id)
+            self.image_end_tag = self.tokenizer.convert_ids_to_tokens(self.config.video_config.image_end_token_id)
+            self.image_pad_tag = self.tokenizer.convert_ids_to_tokens(self.config.video_config.image_pad_token_id)
+            self.video_place_tag = self.tokenizer.convert_ids_to_tokens(self.config.video_config.video_place_token_id)
+
+
+    # @lru_cache(maxsize=1024)
+    def _get_audio(self, audio_info, return_mm_data = True):
+        try:
+            audio_info = ujson.loads(audio_info)
+            audio_uri = None
+            if 'path' in audio_info.keys():
+                
+                if self.relative_path is not None: # 优先匹配本地路径
+                    audio_uri = os.path.join(self.relative_path, audio_info['path'])
+                    if not os.path.exists(audio_uri):  
+                        audio_uri = None
+                if audio_uri is None:  # 本地没有尝试取cos/tos
+                    if audio_info.get('server', 'cos') == 'tos':
+                        audio_uri = self.tos_client(audio_info['path'], 'audio-dataset')
+                    else:
+                        audio_uri = self.cos_client(audio_info['path'], 'audio-data-1317568651')
+
+            elif 'local' in audio_info.keys():
+                audio_uri = audio_info['local']
+                if not os.path.exists(audio_uri):  
+                    audio_uri = None 
+                    return OceanProcessorOutput()                 
+            else:
+                raise ValueError("can not find path or local in audio_info")
+            
+            waveforms = self.audio_processor.load_audio_waveform(audio_uri, True)
+            waveforms = self.audio_processor.split_with_overlap(waveforms)  # 分割逻辑
+            ret = OceanProcessorOutput()  # 默认初始化 audios字段为None
+            for waveform in waveforms:
+                audio, input_length = self.audio_processor.extract_fbank_features(waveform)
+                audio = self.audio_processor.data_augment(audio, input_length, self.training)
+                encoder_length, bridge_length = self.audio_processor.inference_output_length(self.config.audio_config, input_length)
+                if bridge_length <= 0:  # 过滤极端短数据 1. 如果len(waveforms)==1 ret=None; 2. len(waveforms)>1 则说明最后一段太短被抛弃
+                    continue
+                current_ret = OceanProcessorOutput(
+                    audios=[audio], 
+                    encoder_length=[encoder_length], 
+                    bridge_length=[bridge_length])
+                if ret.audios is None:
+                    ret = current_ret
+                else:
+                    ret = ret.concatenate(current_ret)  # 拼接多个切片
+            if not return_mm_data:
+                ret.audios = [None]
+            return ret
+        except Exception as e:
+            print("**** get audio error: {}, info: {} *****".format(str(e), str(audio_info)))
+        return OceanProcessorOutput()
+
+    # @lru_cache(maxsize=1024)
+    def _get_image(self, image_info, return_mm_data = True):
+        try:
+            try: # chensong
+                image_info = ujson.loads(image_info)
+            except:
+                #image_info = image_info.replace("'", '"')
+                image_info = re.sub(r"(?<!\\)'", '"', image_info)
+                image_info = ujson.loads(image_info)
+            if 'base64' in image_info.keys():
+                image_data = base64.b64decode(image_info['base64'])
+                image_feat, org_size, image_list = self.visual_processor.image_transform(image_data)
+            elif 'local' in image_info.keys():
+                image_feat, org_size, image_list = self.visual_processor.image_transform(image_info['local'],return_mm_data = return_mm_data)
+            elif 'path' in image_info.keys():
+                if "tos_bucket" in image_info.keys(): # tos上的每个item，一定要写明tos的桶以及tos_bucket这个key
+                    tos_bucket = image_info['tos_bucket']
+                    image_bytes = self.tos_client(image_info['path'], tos_bucket) # 从cos_client 获得 image
+                else:
+                    cos_bucket = None
+                    if "cos_bucket" in image_info.keys():
+                        cos_bucket = image_info['cos_bucket']
+                    if "bucket_name" in image_info.keys():
+                        cos_bucket = image_info['bucket_name']
+                    image_bytes = self.cos_client(image_info['path'], cos_bucket) # 从cos_client 获得 image
+                # 获得image_feat(image patches), org_size(image最初的size), image_list
+                image_feat, org_size, image_list = self.visual_processor.image_transform(image_bytes)
+            else:
+                raise ValueError("can not find any path in image_info")
+            
+            merge_length = self.visual_processor.merge_size**2
+            patch_nums = np.array(image_list).prod() // merge_length
+            
+            if org_size[0] * org_size[1] > 16**2:  # 极端小的图过滤
+                return OceanProcessorOutput(
+                        images=[image_feat],
+                        patch_nums=[patch_nums],
+                        crop_size=[image_list],
+                        images_size= [org_size],
+                        images_grid=[image_list]
+                        )
+            else:
+                print("**** image too small: {}, info: {} *****".format(str(org_size), str(image_info)))
+                return OceanProcessorOutput()
+           
+        except Exception as e:
+            print("**** get image error: {}, info: {} *****".format(str(e), str(image_info)))
+        return OceanProcessorOutput()
+    
+    # @lru_cache(maxsize=1024)
+    def _get_video_frame(self, video_frame_info, return_mm_data = True):
+        try:
+            pattern = r'\{.*?\}'
+            matches = re.findall(pattern, video_frame_info)
+            ret = OceanProcessorOutput()
+            # 逐个解析
+            for match in matches:
+                video_frame_info = ujson.loads(match)
+                if 'local' in video_frame_info.keys():
+                    image_feat, org_size, image_list = self.video_processor.image_transform(video_frame_info['local'],return_mm_data = return_mm_data)
+                else:
+                    raise ValueError("can not find any path in image_info")
+                
+                merge_length = self.video_processor.merge_size**2
+                patch_nums = np.array(image_list).prod() // merge_length
+                
+                if org_size[0] * org_size[1] > 16**2:  # 极端小的图过滤
+                    ret = ret.concatenate(
+                            OceanProcessorOutput(
+                                videos=[image_feat],
+                                videos_patch_nums=[patch_nums],
+                                videos_crop_size=[image_list],
+                                videos_size= [org_size],
+                                videos_grid=[image_list]
+                            )
+                        )
+                else:
+                    print("**** video too small: {}, info: {} *****".format(str(org_size), str(video_frame_info)))
+            return ret
+           
+        except Exception as e:
+            print("**** get video error: {}, info: {} *****".format(str(e), str(video_frame_info)))
+        return OceanProcessorOutput()
+    
+    # 读取视频
+    def _get_video_obj_byte(self, source, path, video_obj_json):
+        video_obj_byte = None
+        if source == "cos":
+            start_time = time.time()
+            video_obj_byte = self.cos_client(path, bucket_name=video_obj_json.get("cos_bucket", None))
+            if (time.time() - start_time) > 1.0:
+                self.reflash_cos_client()
+        if source == "local":
+            if os.path.exists(path):
+                video_obj_byte = open(path, "rb").read()
+            else:
+                video_obj_byte = None
+        if source == "base64":
+            video_obj_byte = base64.b64decode(path)
+        if source == "url":
+            video_obj_byte = requests.get(url=path).content
+        return video_obj_byte
+    
+    # 将视频切分为帧，保存至子目录中
+    def _split_video_to_frames(self, video_info, max_frame_number=-1, decode_way="1fps"):
+        video_path = video_info['local']
+        # 帧保存本地路径
+        frame_path = video_path.split('.')[0] + '_frames'
+        if not os.path.exists(frame_path) or len(os.listdir(frame_path))==0:
+            # 保存帧
+            os.makedirs(frame_path, exist_ok=True)
+            mm_obj_byte = self._get_video_obj_byte('local', video_path, video_info)
+            if mm_obj_byte is None: # 未读取到视频文件
+                return ""
+            frames = read_video(io.BytesIO(mm_obj_byte), max_frame_number=max_frame_number, decode_way=decode_way) #读取全部帧
+            for frame_idx, frame in enumerate(frames):
+                output_filename = os.path.join(frame_path, f"{frame_idx}.jpg")
+                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+                cv2.imwrite(output_filename, frame)
+                    
+        # 选取帧
+        frame_number = len([filename for filename in os.listdir(frame_path) if filename.endswith('.jpg')])
+        if frame_number>max_frame_number:
+            indices = np.linspace(0, frame_number - 1, max_frame_number, dtype=int)
+        else:
+            indices = np.linspace(0, frame_number - 1, frame_number, dtype=int)
+        # 拼接模式
+        replace_str = ""
+        for idx in indices:
+            frame_str = f"{self.image_start_tag}{os.path.join(frame_path, f'{idx}.jpg')}{self.image_end_tag}"
+            replace_str += frame_str
+        return replace_str
+
+    def _get_video_frame_str(self, video_info, return_mm_data = True ):
+        try:
+            video_info = ujson.loads(video_info)
+            if 'local' in video_info.keys():
+                # 获取包含多帧图像路径的字符串，最大帧数量max_frame_number
+                frames_str = self._split_video_to_frames(video_info, max_frame_number=self.config.video_config.max_frame_num, decode_way=self.config.video_config.decode_way)
+                if frames_str != "":
+                    parts = frames_str.split(self.image_end_tag)
+                    result = []
+                    for part in parts:
+                        if self.image_start_tag in part:
+                            before_path, path = part.split(self.image_start_tag)
+                            new_path = f'{self.image_start_tag}{{"local": "{path}"}}{self.image_end_tag}'
+                            result.append(before_path + new_path)
+                        else:
+                            result.append(part)
+                    return ''.join(result)
+            else:
+                raise ValueError('can not find localpath in video_info')
+        except Exception as e:
+            print("**** get video error: {}, info: {} *****".format(str(e), str(video_info)))
+        return ""
+    
+    # def _replace_audio(self, audio_text, return_mm_data = True):
+        # audio_info = re.sub(re.compile(self.audio_start_tag + "|" + self.audio_end_tag), '', audio_text)
+        # ret = self._get_audio(audio_info, return_mm_data)  # 重复取结果 cached result 
+    def _replace_audio(self, audio_text, mminfo_ret_dict):
+        audio_info = re.sub(re.compile(self.audio_start_tag + "|" + self.audio_end_tag), '', audio_text)
+        # ret = self._get_audio(audio_info)  # 重复取结果 cached result    
+        ret = mminfo_ret_dict.get(audio_info, OceanProcessorOutput())    # 直接从字典取
+        if ret.bridge_length is not None:  # TODO 如果pad token很多 tokenizer效率会很低
+            replaced_text = [self.audio_pad_tag * l for l in ret.bridge_length]
+            replaced_text = self.audio_delim_tag.join(replaced_text) 
+            return self.audio_start_tag + replaced_text + self.audio_end_tag
+        return ''
+
+    # def _replace_image(self, image_text, return_mm_data = True):
+    #     image_info = re.sub(re.compile(self.image_start_tag + "|" + self.image_end_tag), '', image_text)
+    #     ret = self._get_image(image_info, return_mm_data)  # 重复取结果 cached result
+    def _replace_image(self, image_text, mminfo_ret_dict):
+        image_info = re.sub(re.compile(self.image_start_tag + "|" + self.image_end_tag), '', image_text)
+        # ret = self._get_image(image_info)  # 重复取结果 cached result
+        ret = mminfo_ret_dict.get(image_info, OceanProcessorOutput())    # 直接从字典取
+        if ret.patch_nums is None:
+            return ''
+        return self.image_start_tag + self.image_pad_tag * ret.patch_nums[0] + self.image_end_tag
+        return ''
+    
+    # def _replace_video_frame(self, video_frame_text, return_mm_data = True):
+        # video_frame_info = re.sub(re.compile(self.image_start_tag + "|" + self.image_end_tag), '', video_frame_text)
+        # ret = self._get_video_frame(video_frame_info, return_mm_data)  # 重复取结果 cached result
+    def _replace_video_frame(self, video_frame_text, mminfo_ret_dict):
+        video_frame_info = re.sub(re.compile(self.video_start_tag + '|' + self.video_end_tag), '', video_frame_text)
+        video_frame_info = re.sub(re.compile(self.image_start_tag + "|" + self.image_end_tag), '', video_frame_info)
+        # ret = self._get_video_frame(video_frame_info)  # 重复取结果 cached result   
+        ret = mminfo_ret_dict.get(video_frame_info, OceanProcessorOutput())
+        if ret.videos_patch_nums is None:
+            return ''
+        video_frame_str = [self.image_start_tag + self.video_place_tag * ret.videos_patch_nums[i] + self.image_end_tag for i in range(len(ret.videos_patch_nums))]
+        return ''.join(video_frame_str)
+        
+    def extract_replace_multimodal(self, text, mtype='audio', return_mm_data = True):
+        # 抽取text中的json格式音频/图像信息，读取并转化为特征，同时估计encoder token数，填入对应数量的pad token
+        if (self.audio_start_tag != None) and (mtype == 'audio'):
+            match_regex = re.compile(self.audio_start_tag + '.*?' + self.audio_end_tag)
+            drop_regex = re.compile(self.audio_start_tag + "|" + self.audio_end_tag)
+            extract_func = self._get_audio
+            replace_func = self._replace_audio
+        elif (self.image_start_tag != None) and (mtype == 'image'):
+            match_regex = re.compile(self.image_start_tag + '.*?' + self.image_end_tag)
+            drop_regex = re.compile(self.image_start_tag + "|" + self.image_end_tag)
+            extract_func = self._get_image
+            replace_func = self._replace_image
+        elif (self.video_start_tag != None) and (mtype == 'video'):
+            video_match_regex = re.compile(self.video_start_tag + '.*?' + self.video_end_tag)
+            video_drop_regex = re.compile(self.video_start_tag + "|" + self.video_end_tag)
+            # 处理视频，将视频路径转换为多帧图像路径
+            mm_info_list = re.findall(video_match_regex, text)
+            for mm_info in mm_info_list:
+                frame_str = self._get_video_frame_str(re.sub(video_drop_regex, '', mm_info))
+                # 替换路径；如果视频不存在，路径替换为空字符串
+                text = re.sub(mm_info, self.video_start_tag + frame_str + self.video_end_tag, text)
+            # 采用多图像处理方式
+            match_regex = re.compile(self.video_start_tag+r'(.*?)'+self.video_end_tag)
+            drop_regex = re.compile(self.image_start_tag + "|" + self.image_end_tag)
+            extract_func = self._get_video_frame
+            replace_func = self._replace_video_frame
+        else:
+            raise ValueError("mtype not supportted!")
+
+        mm_info_list = re.findall(match_regex, text)
+        mm_info_list = [re.sub(drop_regex, '', mm_info) for mm_info in mm_info_list]
+
+        mminfo_ret_dict = {}
+        ret = OceanProcessorOutput()
+        for mm_info in mm_info_list:  # 如果没有匹配到对应的模态 直接返回raw_text=text 结果不会是None
+            mm_ret = extract_func(mm_info, return_mm_data = return_mm_data)
+            mminfo_ret_dict[mm_info] = mm_ret
+            if mm_ret.audios is None and mm_ret.images is None and mm_ret.videos is None:  # 数据包含音频/图像/视频但抽取失败 整条数据无效（ret的raw_text为None
+                return ret
+            ret = ret.concatenate(mm_ret)  # 可能有多条结果，初步collect
+        # ret.raw_text = re.sub(match_regex, lambda x: replace_func(x.group()), text)
+        ret.raw_text = re.sub(match_regex, lambda x: replace_func(x.group(), mminfo_ret_dict), text)
+        return ret
+
+    def process_one(self, text, index=0, raw_only=False, return_mm_data = True):
+        ret = OceanProcessorOutput(index=index)
+        for mtype in self.config.multimodal:  # 循环获取音频 图像结果 并更新raw_text字段
+            mret = self.extract_replace_multimodal(text, mtype, return_mm_data = return_mm_data) # 增加获取视频结果
+            if mret.raw_text is None:  # 数据包含音频但音频获取失败
+                return OceanProcessorOutput(index=index)
+            ret = ret.concatenate(mret)
+            text = mret.raw_text
+            ret.raw_text = text
+        if raw_only:
+            return ret  # 兼容SFT等自定义tokenizer逻辑的代码
+            
+        # 处理预训练中的trainable部分
+        input_ids, labels = [], []
+        trainable_sep = re.findall(r'<trainable_start>|<trainable_end>', ret.raw_text.replace('\n', '<LF>'))
+        if len(trainable_sep) <= 0:
+            input_ids = self.tokenizer(ret.raw_text, padding='do_not_pad', truncation=True, return_tensors="np")['input_ids'][0].tolist()
+            labels = [True for _ in input_ids]
+        else:
+            split_content = re.split(r'<trainable_start>|<trainable_end>', ret.raw_text)
+            for i, sc in enumerate(split_content):
+                if len(sc.strip()) == 0:
+                    continue  # 把多余的空格干掉
+                sc_ids = self.tokenizer(sc, padding='do_not_pad', truncation=True, return_tensors="np")['input_ids'][0].tolist()
+                input_ids.extend(sc_ids)
+                if i == 0 or trainable_sep[i - 1] == '<trainable_end>':  # stop gradient
+                    labels.extend([False] * len(sc_ids))
+                else:
+                    labels.extend([True] * len(sc_ids))
+        # input_ids += [self.tokenizer.eos_token_id]
+        # labels += [True]
+        ret.labels = [input_ids[j] if (l and input_ids[j] not in self.config.multimodal_special_token_no_loss_list) else -100 for j, l in enumerate(labels)]
+        ret.input_ids = input_ids
+        ret.index = index
+        return ret
+
+    @torch.no_grad()
+    def __call__(self, example, parallel=8):
+        # 最终入口 支持预训练数据string，sft数据message， 以及 batch推理数据listofstring 3种形式
+        if isinstance(example, Dict):
+            pass 
+        elif isinstance(example, str):
+            return self.process_one(example)
+        elif isinstance(example, List):  # batch推理 异步多线程处理
+            with cf.ThreadPoolExecutor(min(parallel, len(example))) as executor:
+                future_list = [executor.submit(self.process_one, di, idx) for idx, di in enumerate(example)]
+                batch_data = [key.result() for key in cf.as_completed(future_list)]
+            valid_num = sum([1 if x.input_ids is not None else 0 for x in batch_data])
+            assert(valid_num == len(batch_data))  # 推理数据严格要求数量对齐
+            batch_data = sorted(batch_data, key=lambda x: x.index)  # 保证顺序不变
+            
+            ret = OceanProcessorOutput()
+            for i in range(len(batch_data)):
+                ret = ret.concatenate(batch_data[i])
+            self.tokenizer.padding_side = "left"
+            padding_result = self.tokenizer.pad({"input_ids": [r.input_ids for r in batch_data]}, return_tensors='pt')
+            ret.input_ids = padding_result["input_ids"]
+            ret.attention_mask = padding_result["attention_mask"]  # batch推理不pack 不需要seqlens
+            padding_result = self.tokenizer.pad({"input_ids": [r.labels for r in batch_data]}, return_tensors='pt')
+            ret.labels = padding_result["input_ids"]
+            
+            if ret.audios is not None:
+                ret.audios = default_collate(ret.audios)
+                ret.encoder_length = default_collate(ret.encoder_length)
+                ret.bridge_length = default_collate(ret.bridge_length)
+            
+            if ret.images is not None:
+                ret.images = [torch.from_numpy(np.asarray(image, dtype=np.float32)) for image in ret.images]
+                # else:ret.images = default_collate(ret.images)
+                # ret.patch_nums = default_collate(ret.patch_nums)
+            
+            if ret.videos is not None:
+                ret.images = [torch.from_numpy(np.asarray(image, dtype=np.float32)) for image in ret.videos]
+
+            return ret
+
+        else:
+            raise ValueError("example format supported yet")
+
+    @torch.no_grad()
+    def pack_batch_pretrain(self, raw_batch, max_sequence_length=None, parallel=8):
+        if max_sequence_length is None:
+            max_sequence_length = self.tokenizer.model_max_length
+        # 将N条数据pack为M条 max_sequence_length长度的数据, 每条数据包含所属的多模态输入
+        assert isinstance(raw_batch, List)
+        start_ts = time.time()
+        if parallel > 1:
+            with cf.ThreadPoolExecutor(max_workers=parallel) as executor:
+                future_list = []
+                for idx, json_text in enumerate(raw_batch):
+                    try:  # 读取json
+                        json_obj = ujson.loads(json_text.strip())
+                    except:
+                        try: 
+                            json_obj = ast.literal_eval(json_text.strip())
+                        except:
+                            print("parse json obj faild: {}....".format(json_text[:300]))
+                            continue
+                    try: # chensong
+                        if isinstance(json_obj, list):
+                            content = json_obj[1]
+                        elif 'raw' in json_obj.keys():
+                            content = (json_obj["title"] if "title" in json_obj.keys() else "") + json_obj["raw"]
+                        else:
+                            content = (json_obj["title"] if "title" in json_obj.keys() else "") + json_obj["content"]
+                    except:
+                        print("parse json raw/content error: {}....".format(json_text[:300]))
+                        continue
+     
+                    future_list.append(executor.submit(self.process_one, content, idx))
+                # 获取结果 乱序
+                batch_data = [key.result() for key in cf.as_completed(future_list)]
+        else: # debug only
+            
+            batch_data = []
+            for json_text in raw_batch:
+                data = ujson.loads(json_text.strip())
+                if 'raw' in data.keys():
+                    batch_data.append(self.process_one(data['raw'], 0))
+                else:
+                    batch_data.append(self.process_one(data['content'], 0))
+
+        if (time.time() - start_ts) / (len(batch_data) + 1e-3) > 1.0:
+            print('[WARNING] processing each data cost more than 1.0s')
+
+        # packing 文本部分的输入，不做任何截断
+        current_length, packed_output, output = 0, OceanProcessorOutput(position_ids=[], seqlens=[]), []
+        empty_data = OceanProcessorOutput(input_ids=[], labels=[])
+        for idx, bd in enumerate(batch_data + [empty_data]):  # 加空数据方便appedn最后一个数据到output，防止遗漏
+            if bd.input_ids is None and idx < len(batch_data):
+                continue  # 数据没取到 并且不是最后一个
+            if (len(bd.input_ids) <= 0 or len(bd.input_ids) + 1 > max_sequence_length) and idx < len(batch_data):
+                continue  # 太长的直接不要 并且不是最后一个
+            if current_length + len(bd.input_ids) + 1 > max_sequence_length or idx == len(batch_data):
+                pad_nums = max_sequence_length - current_length  # right padding
+                if packed_output.input_ids is None or packed_output.labels is None:
+                    packed_output.input_ids = [self.tokenizer.pad_token_id] * pad_nums
+                    packed_output.labels = [-100] * pad_nums
+                    packed_output.position_ids += [0] * (pad_nums+1)
+                else:
+                    packed_output.input_ids += [self.tokenizer.pad_token_id] * pad_nums
+                    packed_output.labels += [-100] * pad_nums
+                    packed_output.position_ids += [0] * pad_nums
+                packed_output.attention_mask = [1] * current_length + [0] * pad_nums
+                packed_output.seqlens += [0] * (max_sequence_length - len(packed_output.seqlens))
+                output.append(packed_output)
+                packed_output = OceanProcessorOutput(position_ids=[], seqlens=[])  # reset empty
+            packed_output = packed_output.concatenate(bd)
+            packed_output.input_ids.append(self.tokenizer.eos_token_id)  # </s>需要单独加
+            packed_output.labels.append(self.tokenizer.eos_token_id)
+            
+            packed_output.position_ids.extend(list(range(len(bd.input_ids) + 1)))
+            packed_output.seqlens.append(len(bd.input_ids) + 1)
+
+            current_length = len(packed_output.input_ids)
+        return output
+        
+    @torch.no_grad()
+    def collect_batch_pretrain(self, batch_data):
+        ret = OceanProcessorOutput()
+        for i in range(len(batch_data)):
+            ret = ret.concatenate(batch_data[i])
+        ret.input_ids = default_collate([np.asarray(x.input_ids, dtype=np.int64) for x in batch_data]).cuda(non_blocking=True)
+        ret.labels = default_collate([np.asarray(x.labels, dtype=np.int64) for x in batch_data]).cuda(non_blocking=True)
+        ret.attention_mask = default_collate([np.asarray(x.attention_mask, dtype=np.float32) for x in batch_data]).cuda(non_blocking=True)
+        ret.position_ids = default_collate([np.asarray(x.position_ids, dtype=np.int64) for x in batch_data]).cuda(non_blocking=True)
+        ret.seqlens = default_collate([np.asarray(x.seqlens, dtype=np.int64) for x in batch_data]).cuda(non_blocking=True)
+
+        ret.raw_text = None
+        if ret.audios is not None:
+            ret.audios = default_collate(np.asarray(ret.audios, dtype=np.float32)).cuda(non_blocking=True)
+            ret.encoder_length = default_collate(np.asarray(ret.encoder_length, dtype=np.int32)).cuda(non_blocking=True)
+            ret.bridge_length = default_collate(np.asarray(ret.bridge_length, dtype=np.int32)).cuda(non_blocking=True)
+        if ret.images is not None:
+            ret.images = [torch.from_numpy(np.asarray(image, dtype=np.float32)).cuda(non_blocking=True)  for image in ret.images]#default_collate(np.asarray(ret.images, dtype=np.float32)).cuda(non_blocking=True)
+            ret.patch_nums = default_collate(np.asarray(ret.patch_nums, dtype=np.int32)).cuda(non_blocking=True)
+        if ret.videos is not None:
+            ret.videos = [torch.from_numpy(np.asarray(video, dtype=np.float32)).cuda(non_blocking=True)  for video in ret.videos]#default_collate(np.asarray(ret.images, dtype=np.float32)).cuda(non_blocking=True)
+            ret.videos_patch_nums = default_collate(np.asarray(ret.videos_patch_nums, dtype=np.int32)).cuda(non_blocking=True)
+        
+        return ret
+    
+    @torch.no_grad()
+    def collect_batch_sft(self, batch_data):
+        # list of dict to dataclass
+        batch_data = [OceanProcessorOutput(**bd) for bd in batch_data]
+        ret = OceanProcessorOutput()
+        for i in range(len(batch_data)):
+            ret = ret.concatenate(batch_data[i])
+        ret.input_ids = default_collate([np.asarray(x.input_ids, dtype=np.int64) for x in batch_data])
+        ret.labels = default_collate([np.asarray(x.labels, dtype=np.int64) for x in batch_data])
+        ret.position_ids = default_collate([np.asarray(x.position_ids, dtype=np.int64) for x in batch_data])
+        ret.seqlens = default_collate([np.asarray(x.seqlens, dtype=np.int64) for x in batch_data])
+
+        ret.raw_text = None
+        if ret.audios is not None:
+            ret.audios = default_collate(np.asarray(ret.audios, dtype=np.float32))
+            ret.encoder_length = default_collate(np.asarray(ret.encoder_length, dtype=np.int32))
+            ret.bridge_length = default_collate(np.asarray(ret.bridge_length, dtype=np.int32))
+        if ret.images is not None:
+            # 转换 每个image 为torch tensor
+            ret.images = [torch.from_numpy(np.asarray(image, dtype=np.float32))  for image in ret.images]#default_collate(np.asarray(ret.images, dtype=np.float32)).cuda(non_blocking=True)
+        if ret.videos is not None:
+            ret.videos = [torch.from_numpy(np.asarray(video, dtype=np.float32))  for video in ret.videos]#default_collate(np.asarray(ret.images, dtype=np.float32)).cuda(non_blocking=True)
+        
+            # ret.patch_nums = default_collate(np.asarray(ret.patch_nums, dtype=np.int32)).cuda(non_blocking=True)
+        
+        ret = ret.__dict__
+        del ret['patch_nums']
+        del ret['images_size']
+        del ret['crop_size']
+        del ret['raw_text']
+        del ret['index']
+        del ret['attention_mask']
+        del ret['videos_patch_nums']
+        del ret['videos_size']
+        del ret['videos_crop_size']
+        return ret
+
+
+#######################################################
+## Unit Test Functions, usage
+## python processor_ocean.py test
+#######################################################
+
+def test_img_processor():
+    from transformers import AutoConfig
+    from transformers.models.clip import CLIPImageProcessor
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    processor = OceanImageProcessor(config.visual_config)
+    offical_processor = CLIPImageProcessor(size=config.visual_config.crop_size, crop_size=config.visual_config.crop_size, 
+                        image_mean=config.visual_config.image_mean, image_std=config.visual_config.image_std,
+                        do_convert_rgb=True)
+    img_files = ['sogou/7a2c8ffc1bc61146b32805c3390f42e2', 'wukong/77c1db1c0e4200d12b478c33ba3a412d', 'wukong/62e9a5c8eb8b0ea8858a34ba3f1a999f', 'wukong/fb9ab4d7c3fe9f54289948fd6a57fc30']
+    cos_client = CosClient()
+    for img_file in img_files:
+        img_bytes = cos_client(img_file)
+        img_rbg = Image.open(io.BytesIO(img_bytes))
+        image, org_size = processor.image_transform(img_bytes)
+        offical_image = offical_processor.preprocess([img_rbg],
+                        do_resize=True, do_center_crop=True, do_rescale=True, do_normalize=True, 
+                        return_tensors='np').data['pixel_values'][0]
+        print('-'*60)
+        print(np.array(img_rbg).shape)
+        print(image.shape)
+        print(offical_image.shape)
+        print(image - offical_image)
+
+def test_audio_processor():
+    from transformers.models.whisper import WhisperFeatureExtractor
+    from transformers import AutoConfig
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    offical_processor = WhisperFeatureExtractor(feature_size=128) 
+    processor = OceanAudioProcessor(config.audio_config)
+    # wave_files = glob.glob('/home/nfs_bc_alignment/sunhaoze/audio-data/openaqa/openaqa-as/audio/*')
+    wave_files = ['/home/nfs_bc_alignment/sunhaoze/sounds/audioset_full/7ZY0U5tfKyQ.flac', '/home/nfs_bc_alignment/sunhaoze/sounds/audioset_full/Osly4Shchs4.flac']
+    for wave_file in wave_files:
+        wave = processor.load_audio_waveform(wave_file, True, False)
+        offical_features = offical_processor(wave[0].numpy(), do_normalize=False)
+        feat = offical_features['input_features'][0]
+        wave, frame_nums = processor.extract_fbank_features(wave)
+        print("="*60)
+        print(feat.shape)
+        print(wave.shape, frame_nums)
+        print('the difference between offical extractor and our implementation: {}'.format(wave_file))
+        print(wave[:, :frame_nums] - feat[:, :frame_nums])
+        print(wave)
+        # print(wave[120:-1, :])
+        # print(feat[120:-1, :wave.shape[1]])
+        zeros_before = np.sum(wave == 0)
+        aug = processor.data_augment(wave, frame_nums)
+        zeros_after = np.sum(aug == 0)
+        print(zeros_before, zeros_after)
+
+
+def test_audio_long():  # 测试超过30秒音频的截断策略
+    from transformers import AutoConfig, AutoTokenizer
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    config.audio_config.split_overlap = 1
+    tokenizer = AutoTokenizer.from_pretrained("./", model_max_length=4096)
+    processor = OceanMMProcessor(tokenizer, config, True)
+    examples = ["<audio_start_ocean>{\"path\": \"panda\/testdata\/podcast_demo_30s\/easy_chat_xianliaohuier_30s\/easy_chat_xianliaohuier-133.mp3\"}<audio_end_ocean>What is the level of noise from the speech?\n<trainable_start>The speech energy\n is medium.<trainable_end>",
+             "what's the sound's energy? \n sound1 <audio_start_ocean>{\"path\": \"panda\/testdata\/podcast_demo_30s\/btrt_talk_heihua_30s\/btrt_talk_heihua-116.mp3\"}<audio_end_ocean> \n sound2 <audio_start_ocean>{\"path\": \"panda\/testdata\/podcast_demo_30s\/btrt_talk_heihua_30s\/btrt_talk_heihua-221.mp3\"}<audio_end_ocean>The speech energy is medium.",
+            ]
+    ret = processor(examples)
+    print(ret)
+    print(torch.sum(ret.input_ids == 151659))
+    print(torch.sum(ret.input_ids == 151674))
+
+def test_processor():
+    from transformers import AutoConfig, AutoTokenizer
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained("./", model_max_length=4096)
+    processor = OceanMMProcessor(tokenizer, config, True, '/home/nfs_bc_alignment/sunhaoze/sounds')
+    examples = ["<audio_start_ocean>{\"path\": \"vggsound\/7DH5fqj8j6Q.flac\"}<audio_end_ocean>What is the level of noise from the speech?\n<trainable_start>The speech energy\n is medium.<trainable_end>",
+             "hello, ocean 你好 百川智能。",
+             "what's the sound's energy? \n <audio_start_ocean>{\"path\": \"iemocap\/Ses01F_script01_3_F022.wav\"}<audio_end_ocean>The speech energy is medium.",
+             "sound1: <audio_start_ocean>{\"path\": \"audioset_full\/9B53NVDNT8U.flac\"}<audio_end_ocean>\n sound2: \n<audio_start_ocean>{\"path\": \"audioset_full\/a2dgzb9GDSQ.flac\"}<audio_end_ocean>How is the speech speed related to the estimated speaker age?\n<trainable_start>The slow speech speed suggests a more deliberate and thoughtful approach often seen in mature individuals.<trainable_end>",
+             "<img_start_ocean>{\"path\": \"sogou\/7351ae4f3fbe58ff0e4cc165cfabb3ed\"}<img_end_ocean>新和记潮汕牛肉火锅的牛肉丸好不好吃 用户评价口味怎么样 常州美食牛肉丸实拍图片 大众点评",
+             "这两个图片有什么关系?图片1<img_start_ocean>{\"path\": \"sogou\/ac91d57ab68335913ed41aa283e76356\"}<img_end_ocean>图片2\n<img_start_ocean>{\"path\": \"sogou\/6ad5e632b74265d9ef689e45936ab1aa\"}<img_end_ocean>",
+             "根据图片和语音给出描述\n图片<img_start_ocean>{\"path\": \"sogou\/32274c1ab28d11f8c490cf7ae15b36f1\"}<img_end_ocean>语音<audio_start_ocean>{\"path\": \"voxceleb2\/id06726_s2lysJWkjus_00169.m4a\"}<audio_end_ocean><trainable_start>这是一只猫<trainable_end>",
+             "这些图片和音频不存在<img_start_ocean>{\"path\": \"soogou\/32274c1ab28d11f8c490cf7ae15b36f1\"}<img_end_ocean>语音<audio_start_ocean>{\"path\": \"voxceleb_1\/id06726_s2lysJWkjus_00169.m4a\"}<audio_end_ocean><trainable_start>这是一只猫<trainable_end>"
+            ]
+    ret = processor(examples[4:-1])
+    print(ret)
+    print(torch.sum(ret.input_ids == 151659))
+    print(torch.sum(ret.input_ids == 151662))
+    try:
+        print(ret.bridge_length)
+        print(ret.patch_nums)
+    except:
+        pass
+    print(torch.sum(ret.attention_mask, dim=1))
+
+
+def test_grounding():
+    from transformers import AutoConfig, AutoTokenizer
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained("./", model_max_length=4096)
+    processor = OceanMMProcessor(tokenizer, config, True, '/home/nfs_bc_alignment/sunhaoze/sounds')
+    examples = ["<img_start_ocean>{\"path\": \"grit\/663423bf2f0884c034bf75279bce9694\"}<img_end_ocean>\nWhere is \"A woman\" ? Answer: <trainable_start>The bounding box is <box_start_ocean>(0.58,0.8),(0.71,1.0)<box_end_ocean><trainable_end>",
+             "hello, ocean 你好 百川智能。",
+             "<img_start_ocean>{\"path\": \"grit\/0e6e3952c584cbac7235940a22514656\"}<img_end_ocean> Generate the caption with grounding: <trainable_start>Photo pour Portrait of <ref_start_ocean>young Asian muslim woman wearing hijab<ref_end_ocean><box_start_ocean>(0.09,0.01),(0.77,1.0)<box_end_ocean> shows regret gesture, hand on her forehead, forget something important, against red background - image libre de droit<trainable_end>",
+             "Recognize the object in the outlined section <img_start_ocean>{\"path\": \"grit\/045823cf6f819670f27aee20af7ae0e6\"}<img_end_ocean> of the picture.<box_start_ocean>(0.07,0.2),(0.91,0.96)<box_end_ocean>\n<trainable_start>Inflatable water trampolines<trainable_end>"
+            ]
+    ret = processor(examples)
+    print(ret)
+    for i, input_ids in enumerate(ret.input_ids):
+        print("="*60)
+        print(ret.labels[i])
+
+def test_pack():
+    from transformers import AutoConfig, AutoTokenizer
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained("./", model_max_length=2048)
+    processor = OceanMMProcessor(tokenizer, config, True, '/home/nfs_bc_alignment/sunhaoze/sounds')
+    examples = open('/cpfs/29f69eb5e2e60f26/user/sunhaoze/pretrain-v6/sogou/part-00000').readlines()[:5]
+    examples += open('/home/nfs_bc_alignment/sunhaoze/text/openaqa-as-stage2-v1/part-00000').readlines()[:5]
+    random.shuffle(examples)
+    batch_output = processor.pack_batch_pretrain(examples)
+    for i, b in enumerate(batch_output):
+        print('='*60)
+        try:
+            print(b.input_ids, len(b.input_ids))
+            print(b.labels, len(b.labels))
+            print(b.attention_mask, len(b.attention_mask))
+            print(b.position_ids, len(b.position_ids))
+            print(b.seqlens, len(b.seqlens))
+            print(b.audios)
+            print(b.bridge_length)
+        except:
+            continue
+    
+    batch_for_model = processor.collect_batch_pretrain(batch_output)
+    print(batch_for_model.input_ids.shape)
+    print(batch_for_model.labels.shape)
+    print(batch_for_model.audios.shape)
+    print(batch_for_model["bridge_length"])
+    print(batch_for_model.images.shape)
+    print(batch_for_model["patch_nums"])
+    print(batch_for_model["position_ids"])
+    print(batch_for_model["seqlens"])
+
+def test_cos_audio():
+    cos_client = CosClient()
+    audio_bytes = cos_client('panda/data/common_voice/cv-corpus-18.0-2024-06-14/zh-CN/clips/common_voice_zh-CN_19428637.mp3', 'audio-data-1317568651')
+    wave, sr = torchaudio.load(audio_bytes, normalize=False)
+    print(wave.shape, sr)
+    # torchaudio.save('tmp.flac', wave, sr)
+
+if __name__ == '__main__':
+    fire.Fire()

special_tokens_map.json

@@ -24,27 +24,27 @@
     "<calc_start>",
     "<calc_end>",
     "<inner_think>",
-    "<audio_start_baichuan>",
-    "<audio_end_baichuan>",
-    "<audio_pad_baichuan>",
-    "<img_start_baichuan>",
-    "<img_end_baichuan>",
-    "<img_pad_baichuan>",
-    "<img_newline_baichuan>",
-    "<box_start_baichuan>",
-    "<box_end_baichuan>",
-    "<box_delim_baichuan>",
-    "<ref_start_baichuan>",
-    "<ref_end_baichuan>",
-    "<img_delim_baichuan>",
-    "<polygon_start_baichuan>",
-    "<polygon_end_baichuan>",
-    "<baichuan_pad_token>",
+    "<audio_start_ocean>",
+    "<audio_end_ocean>",
+    "<audio_pad_ocean>",
+    "<img_start_ocean>",
+    "<img_end_ocean>",
+    "<img_pad_ocean>",
+    "<img_newline_ocean>",
+    "<box_start_ocean>",
+    "<box_end_ocean>",
+    "<box_delim_ocean>",
+    "<ref_start_ocean>",
+    "<ref_end_ocean>",
+    "<img_delim_ocean>",
+    "<polygon_start_ocean>",
+    "<polygon_end_ocean>",
+    "<ocean_pad_token>",
     "<reserved_113>",
-    "<audio_delim_baichuan>",
-    "<video_start_baichuan>",
-    "<video_end_baichuan>",
-    "<video_palce_baichuan>"
+    "<audio_delim_ocean>",
+    "<video_start_ocean>",
+    "<video_end_ocean>",
+    "<video_palce_ocean>"
   ],
   "eos_token": {
     "content": "<|endoftext|>",

tokenizer.json

@@ -302,7 +302,7 @@
     },
     {
       "id": 151676,
-      "content": "<audio_start_baichuan>",
+      "content": "<audio_start_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -311,7 +311,7 @@
     },
     {
       "id": 151677,
-      "content": "<audio_end_baichuan>",
+      "content": "<audio_end_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -320,7 +320,7 @@
     },
     {
       "id": 151678,
-      "content": "<audio_pad_baichuan>",
+      "content": "<audio_pad_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -329,7 +329,7 @@
     },
     {
       "id": 151679,
-      "content": "<img_start_baichuan>",
+      "content": "<img_start_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -338,7 +338,7 @@
     },
     {
       "id": 151680,
-      "content": "<img_end_baichuan>",
+      "content": "<img_end_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -347,7 +347,7 @@
     },
     {
       "id": 151681,
-      "content": "<img_pad_baichuan>",
+      "content": "<img_pad_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -356,7 +356,7 @@
     },
     {
       "id": 151682,
-      "content": "<img_newline_baichuan>",
+      "content": "<img_newline_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -365,7 +365,7 @@
     },
     {
       "id": 151683,
-      "content": "<box_start_baichuan>",
+      "content": "<box_start_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -374,7 +374,7 @@
     },
     {
       "id": 151684,
-      "content": "<box_end_baichuan>",
+      "content": "<box_end_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -383,7 +383,7 @@
     },
     {
       "id": 151685,
-      "content": "<box_delim_baichuan>",
+      "content": "<box_delim_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -392,7 +392,7 @@
     },
     {
       "id": 151686,
-      "content": "<ref_start_baichuan>",
+      "content": "<ref_start_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -401,7 +401,7 @@
     },
     {
       "id": 151687,
-      "content": "<ref_end_baichuan>",
+      "content": "<ref_end_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -410,7 +410,7 @@
     },
     {
       "id": 151688,
-      "content": "<img_delim_baichuan>",
+      "content": "<img_delim_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -419,7 +419,7 @@
     },
     {
       "id": 151689,
-      "content": "<polygon_start_baichuan>",
+      "content": "<polygon_start_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -428,7 +428,7 @@
     },
     {
       "id": 151690,
-      "content": "<polygon_end_baichuan>",
+      "content": "<polygon_end_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -437,7 +437,7 @@
     },
     {
       "id": 151691,
-      "content": "<baichuan_pad_token>",
+      "content": "<ocean_pad_token>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -455,7 +455,7 @@
     },
     {
       "id": 151693,
-      "content": "<audio_delim_baichuan>",
+      "content": "<audio_delim_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -464,7 +464,7 @@
     },
     {
       "id": 151694,
-      "content": "<video_palce_baichuan>",
+      "content": "<video_palce_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -473,7 +473,7 @@
     },
     {
       "id": 151695,
-      "content": "<video_start_baichuan>",
+      "content": "<video_start_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -482,7 +482,7 @@
     },
     {
       "id": 151696,
-      "content": "<video_end_baichuan>",
+      "content": "<video_end_ocean>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,

tokenizer_config.json

@@ -267,7 +267,7 @@
       "special": true
     },
     "151676": {
-      "content": "<audio_start_baichuan>",
+      "content": "<audio_start_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -275,7 +275,7 @@
       "special": true
     },
     "151677": {
-      "content": "<audio_end_baichuan>",
+      "content": "<audio_end_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -283,7 +283,7 @@
       "special": true
     },
     "151678": {
-      "content": "<audio_pad_baichuan>",
+      "content": "<audio_pad_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -291,7 +291,7 @@
       "special": true
     },
     "151679": {
-      "content": "<img_start_baichuan>",
+      "content": "<img_start_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -299,7 +299,7 @@
       "special": true
     },
     "151680": {
-      "content": "<img_end_baichuan>",
+      "content": "<img_end_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -307,7 +307,7 @@
       "special": true
     },
     "151681": {
-      "content": "<img_pad_baichuan>",
+      "content": "<img_pad_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -315,7 +315,7 @@
       "special": true
     },
     "151682": {
-      "content": "<img_newline_baichuan>",
+      "content": "<img_newline_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -323,7 +323,7 @@
       "special": true
     },
     "151683": {
-      "content": "<box_start_baichuan>",
+      "content": "<box_start_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -331,7 +331,7 @@
       "special": true
     },
     "151684": {
-      "content": "<box_end_baichuan>",
+      "content": "<box_end_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -339,7 +339,7 @@
       "special": true
     },
     "151685": {
-      "content": "<box_delim_baichuan>",
+      "content": "<box_delim_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -347,7 +347,7 @@
       "special": true
     },
     "151686": {
-      "content": "<ref_start_baichuan>",
+      "content": "<ref_start_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -355,7 +355,7 @@
       "special": true
     },
     "151687": {
-      "content": "<ref_end_baichuan>",
+      "content": "<ref_end_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -363,7 +363,7 @@
       "special": true
     },
     "151688": {
-      "content": "<img_delim_baichuan>",
+      "content": "<img_delim_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -371,7 +371,7 @@
       "special": true
     },
     "151689": {
-      "content": "<polygon_start_baichuan>",
+      "content": "<polygon_start_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -379,7 +379,7 @@
       "special": true
     },
     "151690": {
-      "content": "<polygon_end_baichuan>",
+      "content": "<polygon_end_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -387,7 +387,7 @@
       "special": true
     },
     "151691": {
-      "content": "<baichuan_pad_token>",
+      "content": "<ocean_pad_token>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -403,7 +403,7 @@
       "special": true
     },
     "151693": {
-      "content": "<audio_delim_baichuan>",
+      "content": "<audio_delim_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -411,7 +411,7 @@
       "special": true
     },
     "151694": {
-      "content": "<video_palce_baichuan>",
+      "content": "<video_palce_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -419,7 +419,7 @@
       "special": true
     },
     "151695": {
-      "content": "<video_start_baichuan>",
+      "content": "<video_start_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -427,7 +427,7 @@
       "special": true
     },
     "151696": {
-      "content": "<video_end_baichuan>",
+      "content": "<video_end_ocean>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -460,27 +460,27 @@
     "<calc_start>",
     "<calc_end>",
     "<inner_think>",
-    "<audio_start_baichuan>",
-    "<audio_end_baichuan>",
-    "<audio_pad_baichuan>",
-    "<img_start_baichuan>",
-    "<img_end_baichuan>",
-    "<img_pad_baichuan>",
-    "<img_newline_baichuan>",
-    "<box_start_baichuan>",
-    "<box_end_baichuan>",
-    "<box_delim_baichuan>",
-    "<ref_start_baichuan>",
-    "<ref_end_baichuan>",
-    "<img_delim_baichuan>",
-    "<polygon_start_baichuan>",
-    "<polygon_end_baichuan>",
-    "<baichuan_pad_token>",
+    "<audio_start_ocean>",
+    "<audio_end_ocean>",
+    "<audio_pad_ocean>",
+    "<img_start_ocean>",
+    "<img_end_ocean>",
+    "<img_pad_ocean>",
+    "<img_newline_ocean>",
+    "<box_start_ocean>",
+    "<box_end_ocean>",
+    "<box_delim_ocean>",
+    "<ref_start_ocean>",
+    "<ref_end_ocean>",
+    "<img_delim_ocean>",
+    "<polygon_start_ocean>",
+    "<polygon_end_ocean>",
+    "<ocean_pad_token>",
     "<reserved_113>",
-    "<audio_delim_baichuan>",
-    "<video_start_baichuan>",
-    "<video_end_baichuan>",
-    "<video_palce_baichuan>"
+    "<audio_delim_ocean>",
+    "<video_start_ocean>",
+    "<video_end_ocean>",
+    "<video_palce_ocean>"
   ],
   "bos_token": null,
   "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",

visual_modeling_ocean.py

@@ -0,0 +1,166 @@
+
+from typing import List, Optional, Tuple, Union
+import torch, math
+import torch.utils.checkpoint
+from torch import nn
+import transformers
+from flash_attn import flash_attn_varlen_func
+from transformers.activations import ACT2FN
+from PIL import Image
+import io, fire
+from torch.nn import functional as F
+
+
+class OceanVisualAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5  # flash attention不需要使用
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        # print("*******OceanVisualAttention monkey patched!!*******")
+
+        # initializer transformer from_pretrain下不生效
+        factor = self.config.initializer_factor
+        in_proj_std = (self.embed_dim**-0.5) * ((2 * self.config.num_hidden_layers) ** -0.5) * factor
+        out_proj_std = (self.embed_dim**-0.5) * factor
+        nn.init.normal_(self.q_proj.weight, std=in_proj_std)
+        nn.init.normal_(self.k_proj.weight, std=in_proj_std)
+        nn.init.normal_(self.v_proj.weight, std=in_proj_std)
+        nn.init.normal_(self.out_proj.weight, std=out_proj_std)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        causal_attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        bsz, tgt_len, embed_dim = hidden_states.size()
+        src_len = tgt_len
+
+        query_states = self.q_proj(hidden_states).view(bsz * tgt_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz * tgt_len, self.num_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz * tgt_len, self.num_heads, self.head_dim)
+
+        # 暂时不考虑变长patch nums 固定长度为256/1024
+        cu_len = torch.arange(0, (bsz + 1) * tgt_len, step=tgt_len, dtype=torch.int32, device=query_states.device)
+        # print(self.config.s2a, self.config.rope_scaling, cu_len, torch.sum(cu_len), q_len, kv_seq_len)
+        # 如果不是f16 bf16不用flash attn
+        if query_states.dtype in [torch.float16, torch.bfloat16]:
+            attn_output = flash_attn_varlen_func(query_states, key_states, value_states, cu_len, cu_len, tgt_len, tgt_len, causal=False)  # (bsz * qlen, nheads, headdim)
+            attn_output = attn_output.view(bsz, tgt_len, self.num_heads, self.head_dim)
+        else:
+            with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
+                attn_output = torch.nn.functional.scaled_dot_product_attention(query_states, key_states, value_states, attention_mask, 0.0)
+                attn_output = attn_output.transpose(1, 2)    
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None
+
+# monkey patch for flash attention
+# transformers.models.siglip.modeling_siglip.SiglipAttention = OceanVisualAttention
+from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VisionTransformerPretrainedModel
+from qwen_vl_utils import process_vision_info
+
+class OceanVisualEncoder(transformers.models.qwen2_vl.modeling_qwen2_vl.Qwen2VisionTransformerPretrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.gradient_checkpointing = True  # 强制开启
+        self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
+        del self.merger
+        
+    def forward(
+        self,
+        pixel_values: torch.Tensor, 
+        grid_thw: torch.Tensor,
+    ):
+        hidden_states = pixel_values.to(self.get_dtype())
+        grid_thw = grid_thw.to(pixel_values.device)
+        
+        hidden_states = self.patch_embed(hidden_states)
+        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+            dim=0, dtype=torch.int32
+        )
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+        for blk in self.blocks:
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(blk.__call__, hidden_states, cu_seqlens, rotary_pos_emb)
+            else:
+                hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)
+
+        return hidden_states
+    
+    @torch.no_grad()
+    def fake_input(self, device, merge_size=2):
+        merge_size = max(merge_size, self.config.spatial_merge_size)
+        fake_image = [torch.zeros([
+            1,
+            self.config.temporal_patch_size,
+            3,
+            merge_size // self.config.spatial_merge_size,
+            self.config.spatial_merge_size,
+            self.config.patch_size,
+            merge_size // self.config.spatial_merge_size,
+            self.config.spatial_merge_size,
+            self.config.patch_size,
+        ], dtype=torch.float32, device=device)]
+        return fake_image
+
+
+class OceanVisualBridge(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.merge_size = self.config.merge_size if hasattr(self.config, 'merge_size') else 2
+        self.hidden_size = config.embed_dim * (self.merge_size**2)
+        self.ln_q = nn.LayerNorm(config.embed_dim, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(self.hidden_size, self.hidden_size),
+            nn.GELU(),
+            nn.Linear(self.hidden_size, config.hidden_size),
+        )
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
+        return x
+    
+
+def test_vision():
+    from transformers.models.clip.modeling_clip import CLIPPreTrainedModel
+    from transformers import AutoConfig
+    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
+
+    ae = OceanVisualEncoder(config.visual_config).cuda().to(torch.bfloat16)
+    bg = OceanVisualBridge(config).cuda().to(torch.bfloat16)
+    print(ae)
+    pixel_input = torch.rand([4, 3, config.visual_config.image_size, config.visual_config.image_size], dtype=torch.float32).cuda()
+
+    visual_embedding = ae(pixel_input)[0][:, 1:]  # 删除class token
+    visual_proj = bg(visual_embedding)
+    print(visual_proj.shape)
+    print(ae.fake_input(visual_proj.device))
+
+if __name__ == '__main__':
+    fire.Fire()
+