File size: 5,640 Bytes
0ab2566 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import logging
from dataclasses import dataclass
from typing import Optional, Union
import torch
import torch.nn.functional as F
from transformers.cache_utils import Cache
from transformers.modeling_outputs import ModelOutput
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
logger = logging.getLogger(__name__)
IGNORE_INDEX = -100
@dataclass
class Qwen3VLCoTEmbeddingOutput(ModelOutput):
"""CoT embedding 模型的前向输出。
last_hidden_state / attention_mask 用于 last-token pooling 取 embedding；
lm_loss 为 CoT 段上的 cross entropy，只在传入 labels 时计算；
logits / past_key_values 只在生成（return_logits=True）时填充。
"""
lm_loss: Optional[torch.FloatTensor] = None
num_cot_tokens: Optional[torch.Tensor] = None
last_hidden_state: Optional[torch.FloatTensor] = None
attention_mask: Optional[torch.Tensor] = None
logits: Optional[torch.FloatTensor] = None
past_key_values: Optional[Cache] = None
class Qwen3VLForCoTEmbedding(Qwen3VLForConditionalGeneration):
"""在 Qwen3-VL 上同时输出 CoT 生成 logits 与检索 embedding 的模型。
继承 Qwen3VLForConditionalGeneration 以复用 lm_head、权重 tie 逻辑以及
generate()，因此同一份权重既能做 CoT 生成（推理阶段 query 先出 CoT）
也能做 embedding（对 last_hidden_state 做 last-token pooling）。
query 侧的输入由 collator 拼成 `prefix + cot + <|im_end|>` 的 teacher forcing
形式，labels 只在 CoT 段有效；doc 侧不传 labels，只取 embedding。
"""
accepts_loss_kwargs = False
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Cache] = None,
inputs_embeds: Optional[torch.FloatTensor] = 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,
cache_position: Optional[torch.LongTensor] = None,
labels: Optional[torch.LongTensor] = None,
return_logits: bool = False,
logits_to_keep: Union[int, torch.Tensor] = 0,
**kwargs,
) -> Qwen3VLCoTEmbeddingOutput:
outputs = self.model(
input_ids=input_ids,
pixel_values=pixel_values,
pixel_values_videos=pixel_values_videos,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
position_ids=position_ids,
attention_mask=attention_mask,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
cache_position=cache_position,
**kwargs,
)
hidden_states = outputs.last_hidden_state
lm_loss, num_cot_tokens = None, None
if labels is not None:
lm_loss, num_cot_tokens = self.cot_lm_loss(hidden_states, labels)
# 训练时不需要全序列 logits（显存开销 = B*L*vocab），只有生成时才算
logits = None
if return_logits:
slice_indices = (
slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) and logits_to_keep > 0
else slice(None)
)
logits = self.lm_head(hidden_states[:, slice_indices, :])
return Qwen3VLCoTEmbeddingOutput(
lm_loss=lm_loss,
num_cot_tokens=num_cot_tokens,
last_hidden_state=hidden_states,
attention_mask=attention_mask,
logits=logits,
past_key_values=outputs.past_key_values,
)
def cot_lm_loss(self, hidden_states: torch.Tensor, labels: torch.Tensor):
"""只在 labels != -100 的位置上算 next-token cross entropy。
先按 label mask 把需要预测的 hidden state 挑出来再过 lm_head，
避免对整条序列展开 (B, L, vocab_size) 的 logits。
"""
shift_hidden = hidden_states[:, :-1, :]
shift_labels = labels[:, 1:]
keep = shift_labels != IGNORE_INDEX
num_cot_tokens = keep.sum()
if num_cot_tokens == 0:
# 整个 batch 都没有 CoT 监督（例如全部被截断），返回挂在计算图上的 0
logger.warning("No CoT label token in this batch, cot loss is skipped.")
return hidden_states.sum() * 0.0, num_cot_tokens
selected_hidden = shift_hidden[keep] # (M, H)
targets = shift_labels[keep].to(torch.long) # (M,)
selected_logits = self.lm_head(selected_hidden).float() # (M, vocab_size)
return F.cross_entropy(selected_logits, targets), num_cot_tokens
def prepare_inputs_for_generation(self, *args, **kwargs):
# generate() 依赖 outputs.logits，因此生成路径上显式打开 logits 计算
model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
model_inputs["return_logits"] = True
return model_inputs
def pooling_last(hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
"""取每条序列最后一个有效 token 的 hidden state 作为 embedding。"""
flipped = attention_mask.flip(dims=[1])
last_positions = flipped.argmax(dim=1)
col = attention_mask.shape[1] - last_positions - 1
row = torch.arange(hidden_state.shape[0], device=hidden_state.device)
return hidden_state[row, col]
|