| 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 = 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: |
| |
| 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] |
| targets = shift_labels[keep].to(torch.long) |
| selected_logits = self.lm_head(selected_hidden).float() |
| return F.cross_entropy(selected_logits, targets), num_cot_tokens |
|
|
| def prepare_inputs_for_generation(self, *args, **kwargs): |
| |
| 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] |
|
|
