modeling_valley.py

#  Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# 
#  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.

import torch
import numpy as np
from torch import nn
from torch.nn import CrossEntropyLoss
from abc import ABC, abstractmethod
from typing import List, Optional, Tuple, Union, Dict, Any
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers import AutoConfig, AutoModelForCausalLM, Qwen3Config, Qwen3ForCausalLM, Qwen3Model

from .modeling_vision_tower import build_vision_tower
from .modeling_projector import build_vision_projector
from .utils import get_anyres_image_grid_shape, unpad_image, IGNORE_INDEX, IMAGE_TOKEN_INDEX, IMAGE_INDICATOR_IDS, IMAGE_ATOM_ID, pad_truncate_sequence


class ValleyConfig(Qwen3Config):
    model_type = "valley"

class ValleyMetaModel:
    def __init__(self, config):
        super(ValleyMetaModel, self).__init__(config)
        # Build vision tower
        if hasattr(config, "mm_vision_tower"):
            if getattr(config, "eagle_vision_tower", None) is not None:
                self.vision_tower, self.qwen2vl_vision_tower = build_vision_tower(config, delay_load=False)
            else:
                self.vision_tower = build_vision_tower(config, delay_load=False)
        # Build Projector
        if hasattr(config, "mm_projector_type") and not getattr(config, "only_navit", False):
            self.mm_projector = build_vision_projector(config)

    def get_vision_tower(self):
        vision_tower = getattr(self, "vision_tower", None)
        if getattr(self.config, "eagle_vision_tower", None) is not None:
            qwen2vl_vision_tower = getattr(self, "qwen2vl_vision_tower", None)
            return vision_tower, qwen2vl_vision_tower
        else:
            return vision_tower

class ValleyMetaForCausalLM(ABC):
    @abstractmethod
    def get_model(self):
        pass

    def get_vision_tower(self):
        return self.get_model().get_vision_tower()

    def split_by_instance(self, original_list, split_sizes):
        start = 0
        sub_lists = []
        for size in split_sizes:
            end = start + size
            sub_list = original_list[start:end]
            sub_lists.append([x.to(self.device) for x in sub_list])
            start = end
        return sub_lists
    
    def encode_images_qwen2vl(self, pixel_values = None, grid_thw = None, split_sizes=None):
        _, qwen2vl_vision_tower = self.get_model().get_vision_tower()
        qwen2vl_image_features = qwen2vl_vision_tower(pixel_values, grid_thw)
        qwen2vl_image_split_sizes = torch.prod(grid_thw[:, 1:3]//2, dim=1)
        qwen2vl_image_features = torch.split(qwen2vl_image_features, qwen2vl_image_split_sizes.tolist(), dim=0)
        qwen2vl_image_features = self.split_by_instance(qwen2vl_image_features, split_sizes)
        return qwen2vl_image_features

    def encode_images(self, images = None, split_sizes = None):
        """
        images: (if not anyres) images.shape = [n,3,336,336] , n = number of images + (number of video) * 8
        images: (if anyres) images.shape = [n,3,336,336] , n = number of tiles * number of images
        """
        if getattr(self.config, "eagle_vision_tower", None) is not None:
            siglip_vision_tower, _ = self.get_model().get_vision_tower()
            image_features = siglip_vision_tower(images)
            image_features = self.get_model().mm_projector(image_features)
        else:
            image_features = self.get_model().get_vision_tower()(images)
            image_features = self.get_model().mm_projector(image_features)

        if getattr(self.config,'anyres', False) and getattr(self.config, 'max_vision_token', None) is not None:
            assert split_sizes is not None
            image_features = list(torch.split(image_features, split_sizes, dim=0))
            for i, image_feature in enumerate(image_features):
                hidden_dim = image_feature.shape[-1]
                image_tokens = image_feature.shape[0]*image_feature.shape[1]
                if getattr(self.config, "eagle_vision_tower", None) is not None:
                    pass # the max_vision_token will be processed in the unpad image token part
                else:
                    if image_tokens > self.config.max_vision_token:
                        intput_shape = int((image_feature.shape[1])**0.5)
                        output_shape = int((self.config.max_vision_token/image_feature.shape[0])**0.5)
                        image_feature = image_feature.view(image_feature.shape[0],intput_shape, intput_shape, -1).permute(0,3,1,2)
                        m = nn.AdaptiveAvgPool2d(output_shape) # different from roi pooling, but in square image, it seems the same
                        pooling_feature = m(image_feature).permute(0,2,3,1)
                        image_features[i] = pooling_feature.view(image_feature.shape[0], -1, hidden_dim)
                split_sizes = None # have already split, set the flag 

        if getattr(self.config, 'mm_use_im_start_end', False):
            raise ValueError('mm_use_im_start is not support')
        if split_sizes is not None:
            image_features = torch.split(image_features, split_sizes, dim=0)
        
        return image_features

    def get_padding_method(self):
        right_padding = getattr(self, 'right_padding', None)
        # if right_padding flag is setted, ignore training flag. 
        if right_padding is not None:
            method = 'right' if right_padding else 'left'
        # in the other way, use training flag to determine the padding method.
        method = 'right' if self.training else 'left'

        return method

    def prepare_inputs_labels_for_multimodal(
        self, input_ids, position_ids, attention_mask, past_key_values, labels, images,
        image_sizes, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw):

        vision_tower = self.get_vision_tower()
        if vision_tower is None or images is None or input_ids.shape[1] == 1:
            if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:
                target_shape = past_key_values[-1][-1].shape[-2] + 1
                attention_mask = torch.cat((attention_mask, torch.ones(
                    (attention_mask.shape[0], target_shape - attention_mask.shape[1]),
                    dtype=attention_mask.dtype,
                    device=attention_mask.device
                )), dim=1)
            return input_ids, position_ids, attention_mask, past_key_values, None, labels

        # Step1: Get image embedings
        if type(images) is list or images.ndim == 5:
            # Without slicing the image
            if not getattr(self.config,'anyres', False) and self.config.mm_projector_type != "ovis2_adapter":
                concat_images = torch.cat([image for image in images], dim=0) # to do batch compute
                split_sizes = [image.shape[0] for image in images]
                
                # Get vision tower feature, check whether only use navit firstly
                if getattr(self.config, 'eagle_vision_tower', None) is not None and getattr(self.config, 'only_navit', False):
                    image_features = None
                else:
                    image_features = self.encode_images(concat_images, split_sizes)
                    image_features = [x.to(self.device) for x in image_features]
                
                # Get Eagle features
                if getattr(self.config, 'eagle_vision_tower', None) is not None:
                    if pixel_values is not None:
                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values, image_grid_thw, split_sizes)
                    elif pixel_values_videos is not None:
                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values_videos, video_grid_thw, split_sizes)
                    else:
                        qwen2vl_image_features = None

            # Slicing the image, each image contains some sub_images:
            # images = [
            #   [image1_tiles(n1,3,336,336), image2_tiles(n2,3,336,336), ...],
            #   [image1_tiles(n1,3,336,336), image2_tiles(n2,3,336,336), ...], ...
            # ]
            else:
                split_sizes = [len(image) for image in images]
                # Get Eagle features
                if getattr(self.config, "eagle_vision_tower", None) is not None:
                    if pixel_values is not None:
                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values, image_grid_thw, split_sizes)
                    elif pixel_values_videos is not None:
                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values_videos, video_grid_thw, split_sizes)
                    else:
                        qwen2vl_image_features = None
                
                # Get vision tower feature, check whether only use navit firstly
                if getattr(self.config, 'eagle_vision_tower', None) is not None and getattr(self.config, 'only_navit', False):
                    image_features = None
                else:
                    image_features = []
                    all_concat_images = []
                    all_split_sizes = []
                    for batch_images in images:
                        concat_images = torch.cat([image for image in batch_images], dim=0) # to do batch compute
                        split_sizes = [image.shape[0] for image in batch_images] 
                        all_concat_images.append(concat_images)
                        all_split_sizes.append(split_sizes)
                    all_image_features = self.encode_images(images=torch.cat(all_concat_images, dim=0), split_sizes=sum(all_split_sizes, []))

                    idx = 0
                    for split_sizes in all_split_sizes:
                        batch_image_features = all_image_features[idx:idx+len(split_sizes)]
                        idx += len(split_sizes)
                        if type(batch_image_features[0]) is list:
                            batch_image_features = [torch.cat(x).to(self.device) for x in batch_image_features]
                        else:
                            batch_image_features = [x.view(-1,x.shape[-1]).to(self.device) for x in batch_image_features] # tiles feature need to flatten in token dimention, [n_tiles, T, d] -> [n_tiles * T, d]
                        image_features.append(batch_image_features)

                if getattr(self.config, "eagle_vision_tower", None) is not None and getattr(self.config, 'only_navit', False) == False:
                    # unpad image tokens
                    height = width = self.config.num_patches_per_side
                    new_image_features = []
                    for batch_image_features, batch_image_sizes in zip(image_features, image_sizes):
                        batch_image_features_list = []
                        for cur_image_feature, cur_image_size in zip(batch_image_features, batch_image_sizes):
                            base_image_feature = cur_image_feature[:width*height, :]
                            image_feature = cur_image_feature[width*height:, :]
                            if image_feature.shape[0] != 0:
                                num_patch_width, num_patch_height = get_anyres_image_grid_shape(
                                    cur_image_size,
                                    self.config.grid_pinpoints,
                                    self.config.vit_crop_size
                                )
                                image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1) # (num_patch_H, num_patch_W, H, W, C)
                                image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() # (C, num_patch_H, H, num_patch_W, W)
                                image_feature = image_feature.flatten(1, 2).flatten(2, 3) # (C, num_token_H, num_token_W)
                                image_feature = unpad_image(image_feature, cur_image_size) # (C, num_token_H_unpad, num_token_W_unpad)
                                input_shape = (image_feature.shape[-2], image_feature.shape[-1])
                                subimage_tokens = np.prod(input_shape)
                                
                                # adaptive avg 2d pool for reducing token num
                                max_subimage_tokens = self.config.max_vision_token-width*height
                                if subimage_tokens > max_subimage_tokens:
                                    aspect_ratio = input_shape[0] / input_shape[1]
                                    output_shape = (
                                        int((max_subimage_tokens/aspect_ratio)**0.5*aspect_ratio),
                                        int((max_subimage_tokens/aspect_ratio)**0.5)
                                    )
                                    m = nn.AdaptiveAvgPool2d(output_shape)
                                    image_feature = m(image_feature)
                                image_feature = image_feature.flatten(1, 2).transpose(0, 1)
                                image_feature = torch.cat((base_image_feature, image_feature), dim=0)
                            else:
                                image_feature = cur_image_feature
                            batch_image_features_list.append(image_feature)
                        new_image_features.append(batch_image_features_list)

                    image_features = new_image_features

        else:
            image_features = self.encode_images(images).to(self.device)


        # Step2: Iterate through each sample in the batch, insert image embedings into input_embeds
        #        and filling labels, attention mask at the same time. Finally, get `new_input_embed`,
        #        `new_labels`, new_attention_mask`.
        _labels = labels
        _position_ids = position_ids
        _attention_mask = attention_mask
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
        if position_ids is None:
            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
        if labels is None:
            labels = torch.full_like(input_ids, IGNORE_INDEX)

        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask.bool())]
        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask.bool())]
        attention_mask = [cur_attention_mask[cur_attention_mask.bool()] for cur_attention_mask in attention_mask]
        
        if self.config.mm_projector_type == "ovis2_adapter":  # for ovis2
            # prepare embedding
            visual_vocab_size = self.config.mlp_hidden_dim
            assert visual_vocab_size == 65536
            text_embedding = self.model.get_input_embeddings()
            visual_indicator_embedding = self.model.mm_projector.embedding(
                torch.tensor(
                    list(range(visual_vocab_size - 5, visual_vocab_size)),
                    dtype=torch.long,
                    device=input_ids[0].device 
                )
            ).to(device=input_ids[0].device)

            new_attention_masks = [] 
            new_input_embeds = []
            new_labels = []
            for i, cur_image_features in enumerate(image_features):
                input_id = input_ids[i]
                text_label = labels[i]


                ovis_image_features = []
                for feature in cur_image_features:
                    ovis_image_features.append(feature)
                ovis_image_features = torch.cat(ovis_image_features, dim=0)

                placeholder_token_mask = torch.lt(input_id, 0)
                text_embed = text_embedding(torch.masked_fill(input_id, placeholder_token_mask, 0))


                for j, indicator_id in enumerate(IMAGE_INDICATOR_IDS):
                    text_embed[input_id == indicator_id] = visual_indicator_embedding[j]
                image_atom_positions = torch.where(torch.eq(input_id, IMAGE_ATOM_ID))[0].tolist()

                input_embed_parts = []
                attention_mask_parts = []
                label_parts = []
                prev_image_atom_position = -1
                # assert ovis_image_features.shape[0] == 256*len(image_atom_positions)
                image_token_len = ovis_image_features.shape[0] // len(image_atom_positions)
                # print("image_token_len: ", image_token_len)
                assert image_token_len in [256, 64]
                if len(image_atom_positions) > 0:
                    for index, image_atom_positions in enumerate(image_atom_positions):
                        input_embed_parts.append(text_embed[prev_image_atom_position + 1:image_atom_positions, :])
                        input_embed_parts.append(ovis_image_features[index*image_token_len:index*image_token_len+image_token_len])  # replace 256
                        label_parts.append(text_label[prev_image_atom_position + 1:image_atom_positions])
                        label_parts.append(torch.full((image_token_len,), IGNORE_INDEX, dtype=torch.long, device=input_id.device))  # remain hypothesis
                        attention_mask_parts.append(
                            torch.ones_like(text_label[prev_image_atom_position + 1:image_atom_positions], dtype=torch.bool, device=input_id.device))
                        attention_mask_parts.append(
                            torch.ones(image_token_len, dtype=torch.bool))
                        
                        prev_image_atom_position = image_atom_positions
                    if prev_image_atom_position + 1 < input_id.shape[0]:
                        input_embed_parts.append(text_embed[prev_image_atom_position + 1:, :])
                        label_parts.append(text_label[prev_image_atom_position + 1:])
                        attention_mask_parts.append(    
                            torch.ones_like(text_label[prev_image_atom_position + 1:], dtype=torch.bool))
                        
                    input_embed = torch.cat([part.to(input_id.device) for part in input_embed_parts], dim=0)
                    attention_mask = torch.cat([part.to(input_id.device) for part in attention_mask_parts], dim=0)
                    label = torch.cat([part.to(input_id.device) for part in label_parts], dim=0)

                    new_input_embeds.append(input_embed)
                    new_attention_masks.append(attention_mask)
                    new_labels.append(label)
                else:
                    raise ValueError(
                        "No image token found in the input. Please check the input_ids and image_features.")
            multimodal_max_length = 0
            left_padding = True
            new_input_embeds = pad_truncate_sequence(multimodal_max_length, new_input_embeds, batch_first=True, padding_value=0.0, left_padding=left_padding)
            new_attention_masks = pad_truncate_sequence(multimodal_max_length, new_attention_masks, batch_first=True, padding_value=False, left_padding=left_padding)
            new_labels = pad_truncate_sequence(multimodal_max_length, new_labels, batch_first=True, padding_value=IGNORE_INDEX, left_padding=left_padding)
            return None, None, new_attention_masks, None, new_input_embeds, new_labels

        else:
            new_input_embeds = []
            new_labels = []
            new_attention_mask = []
            
            for batch_idx, cur_input_ids in enumerate(input_ids):
                cur_batch_image_idx = 0
                num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()

                # Step2-1: If this piece of data is pure text, then concat a dummy image to ensure the whole compute graph is same on all device
                if num_images == 0: 
                    if getattr(self.config, "eagle_vision_tower", None) is not None:
                        if getattr(self.config, 'only_navit', False):
                            cur_image_features = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
                        else:
                            siglip_feat = image_features[batch_idx][cur_batch_image_idx]
                            try:
                                qwen2vl_feat = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
                                cur_image_features = torch.cat((siglip_feat, qwen2vl_feat), dim=0)
                            except Exception as e:
                                print(e)
                                print("only siglip feature:", siglip_feat.shape)
                                cur_image_features = siglip_feat
                    else:
                        cur_image_features = image_features[batch_idx][cur_batch_image_idx]
                    cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
                    cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features.squeeze(0)[0:0]], dim=0)
                    new_input_embeds.append(cur_input_embeds)
                    new_labels.append(labels[batch_idx])
                    new_attention_mask.append(attention_mask[batch_idx])
                    cur_batch_image_idx += 1
                    continue
                
                # Step2-2: Split input_ids, labels, attention_mask by IMAGE_TOKEN_INDEX
                cur_input_ids_noim, cur_labels_noim, cur_attention_mask_noim = [], [], []
                cur_labels = labels[batch_idx]
                cur_attention_mask = attention_mask[batch_idx]
                cur_img_attention_mask = [
                    attention_mask[batch_idx][i].item()
                    for i in torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist()
                ]
                image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]] 
                for i in range(len(image_token_indices) - 1):
                    cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])
                    cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])
                    cur_attention_mask_noim.append(cur_attention_mask[image_token_indices[i]+1:image_token_indices[i+1]])
                split_sizes = [x.shape[0] for x in cur_labels_noim]
                cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
                cur_input_embeds_no_im = list(torch.split(cur_input_embeds, split_sizes, dim=0))# get text features

                # Step2-3: Insert image embedings
                cur_new_input_embeds, cur_new_labels, cur_new_attention_mask = [], [], []
                for i in range(num_images + 1): # to add multimodal feature internal the text feature
                    cur_new_input_embeds.append(cur_input_embeds_no_im[i])
                    cur_new_labels.append(cur_labels_noim[i])
                    cur_new_attention_mask.append(cur_attention_mask_noim[i])
                    if i < num_images:
                        if getattr(self.config, "eagle_vision_tower", None) is not None:
                            if getattr(self.config, 'only_navit', False):
                                cur_image_features = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
                            else:
                                siglip_feat = image_features[batch_idx][cur_batch_image_idx]
                                try:
                                    qwen2vl_feat = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
                                    cur_image_features = torch.cat((siglip_feat, qwen2vl_feat), dim=0)
                                except Exception as e:
                                    print(e)
                                    print("only siglip feature:", siglip_feat.shape)
                                    cur_image_features = siglip_feat
                        else:
                            cur_image_features = image_features[batch_idx][cur_batch_image_idx]
                        cur_batch_image_idx += 1
                        cur_new_input_embeds.append(cur_image_features)
                        cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))
                        cur_new_attention_mask.append(torch.full((cur_image_features.shape[0],), True, device=cur_attention_mask.device, dtype=cur_attention_mask.dtype))

                # Step2-4: Concat image embedings and text embedings
                cur_new_input_embeds = torch.cat(cur_new_input_embeds)
                cur_new_labels = torch.cat(cur_new_labels)
                cur_new_attention_mask = torch.cat(cur_new_attention_mask)
                new_input_embeds.append(cur_new_input_embeds)
                new_labels.append(cur_new_labels)
                new_attention_mask.append(cur_new_attention_mask)

            # Step3: Truncate sequences to max length as image embeddings can make the sequence longer
            tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)
            if tokenizer_model_max_length is not None:
                new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
                new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
                new_attention_mask = [x[:tokenizer_model_max_length] for x in new_attention_mask]

            # Step4: Pad and stack input_embeds, labels, attention_mask
            max_len = max(x.shape[0] for x in new_input_embeds)
            batch_size = len(new_input_embeds)
            new_input_embeds_padded = []
            new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
            new_attention_mask_padded = torch.zeros((batch_size, max_len), dtype=new_attention_mask[0].dtype, device=new_attention_mask[0].device)
            position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)

            for i, (cur_new_embed, cur_new_labels, cur_attention_mask) in enumerate(zip(new_input_embeds, new_labels, new_attention_mask)):
                cur_len = cur_new_embed.shape[0]
                if self.get_padding_method() == 'left':
                    new_input_embeds_padded.append(torch.cat((
                        torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),
                        cur_new_embed
                    ), dim=0))
                    if cur_len > 0:
                        new_labels_padded[i, -cur_len:] = cur_new_labels
                        new_attention_mask_padded[i, -cur_len:] = cur_attention_mask
                        position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
                
                else:
                    new_input_embeds_padded.append(torch.cat((
                        cur_new_embed,
                        torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)
                    ), dim=0))
                    if cur_len > 0:
                        new_labels_padded[i, :cur_len] = cur_new_labels
                        new_attention_mask_padded[i, :cur_len] = cur_attention_mask
                        position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)

            new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
            new_labels = new_labels_padded if _labels is not None else None
            new_attention_mask = new_attention_mask_padded if _attention_mask is not None else None
            if _position_ids is None:
                position_ids = None
            
            return None, position_ids, new_attention_mask, past_key_values, new_input_embeds, new_labels


class ValleyQwen3Model(ValleyMetaModel, Qwen3Model):
    config_class = ValleyConfig
    def __init__(self, config: Qwen3Config):
        super(ValleyQwen3Model, self).__init__(config)


class ValleyQwen3ForCausalLM(Qwen3ForCausalLM, ValleyMetaForCausalLM):
    config_class = ValleyConfig

    def __init__(self, config):
        super(Qwen3ForCausalLM, self).__init__(config)
        self.model = ValleyQwen3Model(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_model(self):
        return self.model

    def _update_model_kwargs_for_generation(
        self,
        outputs: CausalLMOutputWithPast,
        model_kwargs: Dict[str, Any],
        is_encoder_decoder: bool = False,
        num_new_tokens: int = 1,
    ) -> Dict[str, Any]:
        new_model_kwargs = super()._update_model_kwargs_for_generation(
            outputs, 
            model_kwargs, 
            is_encoder_decoder, 
            num_new_tokens
        )
        """
        Set model_kwargs["attention_mask"] to the expanded `attention_mask` in
        the `prepare_inputs_labels_for_multimodal` function to ensure the 
        correctness of the generate behavior when `use_cache` is enabled.
        """
        if not is_encoder_decoder:
            if "attention_mask" in new_model_kwargs:
                attention_mask = outputs.attention_mask
                new_model_kwargs["attention_mask"] = torch.cat(
                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
                )
        return new_model_kwargs


    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        images: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
        image_sizes: Optional[List[List[int]]] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = 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

        if inputs_embeds is None:
            (
                input_ids,
                position_ids,
                attention_mask,
                past_key_values,
                inputs_embeds,
                labels
            ) = self.prepare_inputs_labels_for_multimodal(
                input_ids,
                position_ids,
                attention_mask,
                past_key_values,
                labels,
                images,
                image_sizes,
                pixel_values,
                pixel_values_videos,
                image_grid_thw,
                video_grid_thw,
            )

        # 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,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = CrossEntropyLoss(reduction='mean')
            bs = shift_labels.shape[0]
            shift_labels = shift_labels.to(shift_logits.device)
            loss = torch.stack([loss_fct(shift_logits[i], shift_labels[i]) for i in range(bs)])


        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        res =  CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

        res.attention_mask = attention_mask
        return res

    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:]

        # 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}
        else:
            model_inputs = {"input_ids": input_ids}

        # For cases where only navit (e.g., Valley3) is used, the input `images` may be None.
        # To ensure compatibility with the Valley2 codebase, we need to supplement a dummy image.
        image_sizes = kwargs.get("image_sizes", None)
        if kwargs.get("images", None) is not None:
            images = kwargs.get("images")
        else:
            images = [torch.zeros((len(image_sizes_per_sample), 3, 10, 10)) for image_sizes_per_sample in image_sizes]

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "images": images,
                "image_sizes": kwargs.get("image_sizes", None),
                "pixel_values": kwargs.get("pixel_values", None),
                "pixel_values_videos": kwargs.get("pixel_values_videos", None),
                "image_grid_thw": kwargs.get("image_grid_thw", None),
                "video_grid_thw": kwargs.get("video_grid_thw", None),
            }
        )
        return model_inputs

AutoConfig.register("valley", ValleyConfig)
AutoModelForCausalLM.register(ValleyConfig, ValleyQwen3ForCausalLM)