From b28715e6fd6efade236b5e737eddc039780435b1 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Wed, 4 Feb 2026 16:41:17 +0800 Subject: [PATCH] [update] vlm params --- model/model_vlm.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/model/model_vlm.py b/model/model_vlm.py index 3babdb8..776c0d6 100644 --- a/model/model_vlm.py +++ b/model/model_vlm.py @@ -1,12 +1,11 @@ import os - import torch import warnings from .model_minimind import * -from typing import Optional, Tuple, List +from typing import Optional, Tuple, List, Union from torch import nn from transformers import CLIPProcessor, CLIPModel -from typing import List +from transformers.modeling_outputs import MoeCausalLMOutputWithPast warnings.filterwarnings('ignore') @@ -113,10 +112,10 @@ class MiniMindVLM(MiniMindForCausalLM): def forward(self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, use_cache: bool = False, logits_to_keep: Union[int, torch.Tensor] = 0, + labels: Optional[torch.Tensor] = None, pixel_values: Optional[torch.FloatTensor] = None, **args): batch_size, seq_length = input_ids.shape @@ -166,6 +165,5 @@ class MiniMindVLM(MiniMindForCausalLM): shift_labels = labels[..., 1:].contiguous() loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100) - output = CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=presents, hidden_states=hidden_states) - output.aux_loss = aux_loss + output = MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=presents, hidden_states=hidden_states) return output