"""Mage-Flow text encoder (Qwen3-VL, packed varlen conditioning). Vendored from microsoft/Mage (`mage_flow`, MIT) at commit 76bec2bb3818, with a diffusers-convention wrapper appended. Upstream is the reference implementation: the numerics here are its own functions, not a reimplementation. The mandatory content-policy gate upstream runs in ``generate_images`` is not part of this port. Copyright (c) 2026 Microsoft. Licensed under the MIT License. """ from __future__ import annotations import os from collections.abc import Callable from dataclasses import dataclass try: from typing import Unpack except ImportError: from typing_extensions import Unpack import torch from torch import nn from transformers import AutoProcessor, AutoTokenizer, Cache, Qwen3VLForConditionalGeneration from transformers.cache_utils import DynamicCache from transformers.masking_utils import create_causal_mask from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLCausalLMOutputWithPast, apply_rotary_pos_emb, eager_attention_forward, ) from transformers.utils import ModelOutput # =========================================================================== # Custom Qwen3-VL model (customizable forward output) # =========================================================================== import logging logger = logging.getLogger(__name__) """Attention backend shim — switchable between Flash Attention 2 and 4. Exports a single ``flash_attn_varlen_func`` with the FA2 calling convention. The underlying kernel is selected at runtime via ``set_attn_backend(name)`` (default: ``"flash2"``). The selected kernel is resolved lazily on the first call so model-config-driven selection (which happens after this module is imported) takes effect. Modules that previously did ``from flash_attn import flash_attn_varlen_func`` should import from here instead. For the FA4 path, calling-convention differences are normalised: * ``window_size=(-1, -1)`` (FA2 "no window") -> ``(None, None)`` (FA4). * ``block_table`` -> ``page_table``. * FA4's optional ``(out, lse)`` tuple return is unwrapped to ``out``. * ``dropout_p>0`` / ``alibi_slopes`` / ``return_attn_probs`` raise on FA4. """ from typing import Any, Callable _FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"} _FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"} _SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"} _BACKEND: str = "flash2" _RESOLVED_FN: Callable[..., Any] | None = None def _normalize(name: str) -> str: n = name.lower().strip() if n in _FA2_ALIASES: return "flash2" if n in _FA4_ALIASES: return "flash4" if n in _SDPA_ALIASES: return "sdpa" raise ValueError( f"Unknown attention backend {name!r}; expected one of " f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}" ) def set_attn_backend(name: str) -> None: """Select the flash-attn backend used by ``flash_attn_varlen_func``. Safe to call multiple times; clears the cached resolution on change. """ global _BACKEND, _RESOLVED_FN new = _normalize(name) if new != _BACKEND: _RESOLVED_FN = None _BACKEND = new def _resolve_fa2() -> Callable[..., Any]: # Imported by name, not with a plain ``import``: transformers' dynamic-module loader scans this file # and refuses to load it when it sees an import of a package that is not installed — even one inside # a function that the sdpa fallback never reaches. import importlib return importlib.import_module("flash_attn").flash_attn_varlen_func def _resolve_fa4() -> Callable[..., Any]: # 같은 이유로 이름으로 import (정적 스캔이 하드 요구로 보지 않게) import importlib _fa4_fn = importlib.import_module("flash_attn.cute").flash_attn_varlen_func def _fa4_wrapper( q, k, v, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None, dropout_p: float = 0.0, softmax_scale=None, causal: bool = False, window_size=(-1, -1), softcap: float = 0.0, alibi_slopes=None, deterministic: bool = False, return_attn_probs: bool = False, block_table=None, **_unused: Any, ): if dropout_p and dropout_p > 0: raise NotImplementedError("FA4 backend does not support dropout_p>0") if alibi_slopes is not None: raise NotImplementedError("FA4 backend does not support alibi_slopes") if return_attn_probs: raise NotImplementedError("FA4 backend does not support return_attn_probs") win_l, win_r = window_size if win_l == -1: win_l = None if win_r == -1: win_r = None out = _fa4_fn( q, k, v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, softmax_scale=softmax_scale, causal=causal, window_size=(win_l, win_r), softcap=softcap, deterministic=deterministic, page_table=block_table, return_lse=False, ) if isinstance(out, tuple): out = out[0] return out return _fa4_wrapper def _resolve_sdpa() -> Callable[..., Any]: """FA2 varlen → per-sequence torch.SDPA fallback. Use when flash-attn is unavailable (e.g. CUDA 13 has no prebuilt wheel and source build is brittle). Slower than FA2 (one SDPA dispatch per sequence), but functionally equivalent for the dense / causal / no-alibi paths mageflow actually uses. Window / softcap / alibi / paged-attn / return_attn_probs are not supported and will raise. """ import torch import torch.nn.functional as F def _sdpa_wrapper( q, k, v, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None, dropout_p: float = 0.0, softmax_scale=None, causal: bool = False, window_size=(-1, -1), softcap: float = 0.0, alibi_slopes=None, deterministic: bool = False, return_attn_probs: bool = False, block_table=None, **_unused: Any, ): if dropout_p and dropout_p > 0: raise NotImplementedError("SDPA backend does not support dropout_p>0") if alibi_slopes is not None: raise NotImplementedError("SDPA backend does not support alibi_slopes") if return_attn_probs: raise NotImplementedError("SDPA backend does not support return_attn_probs") if softcap and softcap > 0: raise NotImplementedError("SDPA backend does not support softcap") if window_size not in ((-1, -1), (None, None), (0, 0)): raise NotImplementedError( f"SDPA backend does not support sliding window (got {window_size})" ) if block_table is not None: raise NotImplementedError("SDPA backend does not support paged attention") if cu_seqlens_q is None or cu_seqlens_k is None: raise ValueError("SDPA backend requires cu_seqlens_q and cu_seqlens_k") # GQA: FA2 broadcasts k/v across query head groups natively; torch SDPA # does not (the q vs k head-dim mismatch is the AssertionError "tensor # a (32) must match tensor b (8) at non-singleton dimension 1" we'd see # otherwise). Repeat k/v along the head dim to match q before the loop. n_heads_q = q.shape[1] n_heads_kv = k.shape[1] if n_heads_q != n_heads_kv: if n_heads_q % n_heads_kv != 0: raise ValueError( f"SDPA backend GQA expansion requires q heads ({n_heads_q}) " f"to be divisible by k/v heads ({n_heads_kv})" ) repeat = n_heads_q // n_heads_kv k = k.repeat_interleave(repeat, dim=1) v = v.repeat_interleave(repeat, dim=1) # q/k/v: (total_tokens, nheads, head_dim). Dispatch SDPA per sequence, # then concat. Python-level loop is fine since nseq is small (one per # image in the pack) and image-gen latency is dominated by sampling. cu_q = cu_seqlens_q.tolist() cu_k = cu_seqlens_k.tolist() outs = [] for qs, qe, ks, ke in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]): # (s, h, d) → (1, h, s, d) q_i = q[qs:qe].transpose(0, 1).unsqueeze(0) k_i = k[ks:ke].transpose(0, 1).unsqueeze(0) v_i = v[ks:ke].transpose(0, 1).unsqueeze(0) out_i = F.scaled_dot_product_attention( q_i, k_i, v_i, attn_mask=None, dropout_p=0.0, is_causal=causal, scale=softmax_scale, ) # (1, h, s, d) → (s, h, d) outs.append(out_i.squeeze(0).transpose(0, 1)) return torch.cat(outs, dim=0).contiguous() return _sdpa_wrapper def _resolve() -> Callable[..., Any]: global _RESOLVED_FN if _RESOLVED_FN is None: if _BACKEND == "flash4": _RESOLVED_FN = _resolve_fa4() elif _BACKEND == "sdpa": _RESOLVED_FN = _resolve_sdpa() else: try: _RESOLVED_FN = _resolve_fa2() except ImportError: # flash-attn 2 needs sm80+ and a matching build; sdpa is the portable varlen path, so a # missing kernel falls back instead of failing the load. logger.warning("flash-attn 2 is unavailable; using the sdpa attention backend") _RESOLVED_FN = _resolve_sdpa() return _RESOLVED_FN def flash_attn_varlen_func(*args, **kwargs): return _resolve()(*args, **kwargs) __all__ = ["flash_attn_varlen_func", "set_attn_backend"] @dataclass class Qwen3VLModelOutput(ModelOutput): """Flexible output class for custom Qwen3-VL model.""" loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None past_key_values: Cache | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None last_hidden_state: torch.FloatTensor | None = None attentions: tuple[torch.FloatTensor, ...] | None = None rope_deltas: torch.LongTensor | None = None class CustomQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration): """ Custom Qwen3-VL model that allows customizing the forward output. This class inherits from Qwen3VLForConditionalGeneration and provides hooks to customize what is returned from the forward pass. Example usage: ```python model = CustomQwen3VLForConditionalGeneration.from_pretrained( "Qwen/Qwen3-VL-8B-Instruct", attn_implementation="flash_attention_2" # Use flash attention for faster inference ) # Option 1: Use built-in output modes model.set_output_mode("embedding") # Only return last hidden state (default) model.set_output_mode("full") # Return everything model.set_output_mode("logits") # Only return logits # Option 2: Set a custom output processor def my_custom_output(hidden_states, logits, outputs, **kwargs): return {"embeddings": hidden_states, "pooled": hidden_states.mean(dim=1)} model.set_output_processor(my_custom_output) ``` """ # Output mode constants OUTPUT_MODE_FULL = "full" OUTPUT_MODE_EMBEDDING = "embedding" OUTPUT_MODE_LOGITS = "logits" OUTPUT_MODE_HIDDEN = "hidden" def __init__(self, config): super().__init__(config) self._output_mode = self.OUTPUT_MODE_EMBEDDING self._skip_lm_head = True def set_output_mode(self, mode: str): """ Set the output mode for the forward pass. Args: mode: One of: - "full": Return full Qwen3VLCausalLMOutputWithPast - "embedding": Only return last hidden state (skip lm_head) (default) - "logits": Only return logits - "hidden": Return all hidden states """ valid_modes = [ self.OUTPUT_MODE_FULL, self.OUTPUT_MODE_EMBEDDING, self.OUTPUT_MODE_LOGITS, self.OUTPUT_MODE_HIDDEN, ] if mode not in valid_modes: raise ValueError(f"Invalid output mode: {mode}. Must be one of {valid_modes}") self._output_mode = mode self._skip_lm_head = mode == self.OUTPUT_MODE_EMBEDDING def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, pixel_values: torch.Tensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, **kwargs, ) -> Qwen3VLCausalLMOutputWithPast | Qwen3VLModelOutput | dict | torch.Tensor: """ Forward pass with customizable output. Returns different outputs based on the configured output mode or custom processor. """ 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 ) # Get outputs from the base model (Qwen3VLModel) 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, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, **kwargs, ) # Get the last hidden state hidden_states = outputs[0] # This is the last hidden state # Compute logits if not skipping lm_head logits = None if not self._skip_lm_head: slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) # Compute loss if labels are provided loss = None if labels is not None and logits is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs ) # Return based on output mode if self._output_mode == self.OUTPUT_MODE_EMBEDDING: return Qwen3VLModelOutput( last_hidden_state=hidden_states, past_key_values=outputs.past_key_values, attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, ) elif self._output_mode == self.OUTPUT_MODE_LOGITS: return logits elif self._output_mode == self.OUTPUT_MODE_HIDDEN: return Qwen3VLModelOutput( last_hidden_state=hidden_states, hidden_states=outputs.hidden_states, past_key_values=outputs.past_key_values, attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, ) else: # OUTPUT_MODE_FULL return Qwen3VLCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, ) # =========================================================================== # Packing-aware forward patches (cu_seqlens) for the Qwen3-VL text encoder # =========================================================================== def model_forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, # args for deepstack visual_pos_masks: torch.Tensor | None = None, deepstack_visual_embeds: list[torch.Tensor] | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple | BaseModelOutputWithPast: r""" visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*): The mask of the visual positions. deepstack_visual_embeds (`list[torch.Tensor]`, *optional*): The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim). The feature is extracted from the different visual encoder layers, and fed to the decoder hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334). """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") # torch.jit.trace() doesn't support cache objects in the output if use_cache and past_key_values is None and not torch.jit.is_tracing(): past_key_values = DynamicCache(config=self.config) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) # the hard coded `3` is for temporal, height and width. if position_ids is None: position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) elif position_ids.ndim == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) if position_ids.ndim == 3 and position_ids.shape[0] == 4: text_position_ids = position_ids[0] position_ids = position_ids[1:] else: text_position_ids = position_ids[0] if kwargs.get("cu_seqlens") is None: attention_mask = create_causal_mask( config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=text_position_ids, ) hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers for layer_idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=text_position_ids, past_key_values=past_key_values, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = layer_outputs # add visual features to the hidden states of first several layers if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)): hidden_states = self._deepstack_process( hidden_states, visual_pos_masks, deepstack_visual_embeds[layer_idx], ) hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, past_key_values: Cache | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) cu_seqlens = kwargs.get("cu_seqlens", None) if cu_seqlens is None: attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) else: max_seqlen = torch.diff(cu_seqlens).max().item() if cu_seqlens is not None else None query_states = query_states.transpose(1, 2).squeeze(0) key_states = key_states.transpose(1, 2).squeeze(0) value_states = value_states.transpose(1, 2).squeeze(0) attn_output = flash_attn_varlen_func( q=query_states, k=key_states, v=value_states, cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, causal=True, window_size=(-1, -1), softmax_scale=self.head_dim**-0.5, dropout_p=0.0, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None def qwen3_patch_forward(): """Patch the Qwen3-VL text model + attention forwards to support packed varlen (cu_seqlens) inputs used by ``TextEncoder.forward``.""" from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention, Qwen3VLTextModel Qwen3VLTextModel.forward = model_forward Qwen3VLTextAttention.forward = forward # =========================================================================== # TextEncoder wrapper (packed text -> DiT conditioning embeddings) # =========================================================================== _FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"} _FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"} _SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"} def _resolve_hf_attn_impl(attn_type: str) -> str: """Map a project-level attn_type to a HuggingFace ``attn_implementation`` string. ``VF_HF_ATTN_IMPL`` env var, if set, takes precedence (useful for forcing sdpa on machines without flash-attn). For FA4 we additionally probe that the CUTE-DSL kernel is importable and (when available) ask the HF helper to confirm; if not, fall back to sdpa rather than crashing at load time. """ override = os.environ.get("VF_HF_ATTN_IMPL") if override: return override name = attn_type.lower().strip() if name in _FA2_ALIASES: return "flash_attention_2" if name in _FA4_ALIASES: try: import flash_attn.cute # noqa: F401 fa4_importable = True except Exception: fa4_importable = False if fa4_importable: try: from transformers.utils.import_utils import is_flash_attn_4_available if is_flash_attn_4_available(): return "flash_attention_4" except ImportError: return "flash_attention_4" logger.warning( "attn_type=flash4 requested but flash_attn.cute is unavailable; " "falling back to sdpa for HF text encoder." ) return "sdpa" if name in _SDPA_ALIASES: return "sdpa" raise ValueError( f"Unknown attn_type {attn_type!r}; expected one of " f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}" ) SEQ_MULTI_OF = 32 # --------------------------------------------------------------------------- # transformers-version shim + diffusers component wrapper # --------------------------------------------------------------------------- # ``create_causal_mask`` renamed ``input_embeds`` to ``inputs_embeds`` and dropped ``cache_position`` # after the transformers release this code was written against, so the call above is translated here # rather than edited upstream. _upstream_create_causal_mask = create_causal_mask def _create_causal_mask(*args, **kwargs): if "input_embeds" in kwargs: kwargs["inputs_embeds"] = kwargs.pop("input_embeds") if "cache_position" in kwargs and "cache_position" not in _CREATE_CAUSAL_MASK_PARAMS: kwargs.pop("cache_position") return _upstream_create_causal_mask(*args, **kwargs) import inspect # noqa: E402 _CREATE_CAUSAL_MASK_PARAMS = set(inspect.signature(_upstream_create_causal_mask).parameters) create_causal_mask = _create_causal_mask qwen3_patch_forward() class MageFlowTextEncoder(CustomQwen3VLForConditionalGeneration): """Qwen3-VL text encoder with Mage-Flow's packed (varlen) conditioning forward. ``encode_packed`` is upstream's ``TextEncoder.forward`` body: several prompts are concatenated and isolated by ``cu_seqlens`` in one launch, each sequence's leading template tokens are dropped, and the pooled vector is the mean over what remains. """ def encode_packed(self, input_ids, cu_seqlens, drop_idx: int = 0, inputs: dict | None = None) -> dict: seqlens_list = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().tolist() position_ids = torch.cat([torch.arange(_length, device=input_ids.device) for _length in seqlens_list]) forward_kwargs = { "input_ids": input_ids.unsqueeze(0).to(self.device), "cu_seqlens": cu_seqlens, "position_ids": position_ids.unsqueeze(0).to(self.device), "output_hidden_states": False, "max_seqlen": None, } if inputs is not None: for _key in ("pixel_values", "image_grid_thw"): if inputs.get(_key, None) is not None: forward_kwargs[_key] = inputs[_key].to(self.device) with torch.no_grad(): outputs = self(**forward_kwargs) hidden = outputs.last_hidden_state if getattr(outputs, "last_hidden_state", None) is not None \ else outputs.hidden_states[-1] hidden = hidden.squeeze(0) txt_list, vec_list, valid_lengths = list(), list(), list() for _hidden in torch.split(hidden, seqlens_list, dim=0): _valid = _hidden[drop_idx:] txt_list.append(_valid) vec_list.append(_valid.mean(dim=0)) valid_lengths.append(_valid.shape[0]) return { "txt": torch.cat(txt_list, dim=0), "vec": torch.stack(vec_list, dim=0), "txt_seq_lens": torch.tensor(valid_lengths, device=input_ids.device), }