"""M4SEG batch collation and language-model supervision.""" import os from itertools import zip_longest import numpy as np import torch from model.llava import conversation as conversation_lib from model.llava.mm_utils import tokenizer_image_token from tools.utils import DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IMAGE_TOKEN_INDEX, IGNORE_INDEX def custom_collate_fn_prima( batch, tokenizer=None, use_mm_start_end=True, inference=False, local_rank=-1, img_emb_len=32, qformer_tokenizer=None, padding_side="right", append_eos=True, ): def _is_true(value): return str(value).strip().lower() in ("1", "true", "yes", "y", "on") truncate_from_left = _is_true(os.getenv("PRIMA_TRUNCATE_FROM_LEFT", "0")) seg_token_id = None if tokenizer is not None: try: seg_ids = tokenizer("[SEG]", add_special_tokens=False).input_ids if isinstance(seg_ids, list) and len(seg_ids) == 1: seg_token_id = int(seg_ids[0]) except Exception: seg_token_id = None # Initializing lists and counters image_path_list, global_enc_image_list, grounding_enc_image_list = [], [], [] conversation_list, masks_list = [], [] label_list, resize_list, questions_list = [], [], [] selected_labels_list, inferences = [], [] conv_offset_list, img_offset_list, grounding_offset_list = [0], [0], [0] cnt = img_cnt = grounding_cnt = 0 sample_idxs = [] # Iterating through the batch for item in batch: ( sample_idx, image_path, global_enc_image, grounding_enc_image, conversations, masks, label, resize, questions, sampled_classes, ) = item sample_idxs.append(str(sample_idx)) if not image_path: image_path = [] # Text-only sample: normalize to empty lists / None tensors. # NOTE: This must not contain image tokens in its conversations. global_enc_image = None grounding_enc_image = None label = [] masks = [] resize = [] elif not isinstance(image_path, (tuple, list)): image_path = [image_path] global_enc_image = global_enc_image.unsqueeze(0) grounding_enc_image = ( grounding_enc_image.unsqueeze(0) if grounding_enc_image is not None else None ) if label is None: label = [] else: label = [label] if not isinstance(label, (tuple, list)) else label if masks is None: masks = [] else: masks = [masks] if not isinstance(masks, (tuple, list)) else masks if resize is None: resize = [] else: resize = [resize] if not isinstance(resize, (tuple, list)) else resize # Normalize single-image resize tuples to list-of-tuples so zip_longest doesn't unzip them. if resize is not None and isinstance(resize, (tuple, list)): if len(resize) == 2 and all( isinstance(x, (int, np.integer)) for x in resize ): resize = [tuple(resize)] # Avoid accidental iteration over characters. if isinstance(conversations, str): conversations = [conversations] if isinstance(questions, str): questions = [questions] if questions is None: questions = [] if masks is None: masks = [] if label is None: label = [] if resize is None: resize = [] image_path_list.append(image_path) if global_enc_image is not None: global_enc_image_list.append(global_enc_image) if grounding_enc_image is not None: grounding_enc_image_list.append(grounding_enc_image) conversation_list.extend(conversations) masks_list.append( [(mask.float() if mask is not None else None) for mask in masks] ) label_list.append(label) resize_list.append(resize) questions_list.extend(questions) selected_labels_list.append(sampled_classes) conv_offset_list.append(cnt := cnt + len(conversations)) if global_enc_image is not None: img_offset_list.append(img_cnt := img_cnt + global_enc_image.shape[0]) else: img_offset_list.append(img_cnt) if grounding_enc_image is not None: grounding_offset_list.append( grounding_cnt := grounding_cnt + grounding_enc_image.shape[0] ) else: grounding_offset_list.append(grounding_cnt) inferences.append(inference) # Handling the conversation list if use_mm_start_end: replace_token = ( DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN ) conversation_list = [ conv.replace(DEFAULT_IMAGE_TOKEN, replace_token) for conv in conversation_list ] # Tokenizing and padding input ids def _tok(prompt): ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt") if ( (not append_eos) and ids.numel() > 0 and ids[-1].item() == tokenizer.eos_token_id ): ids = ids[:-1] num_images = int((ids == IMAGE_TOKEN_INDEX).sum().item()) truncate_len = tokenizer.model_max_length - (img_emb_len - 1) * num_images truncate_len = max(1, int(truncate_len)) if ids.numel() > truncate_len: ids = ids[-truncate_len:] if truncate_from_left else ids[:truncate_len] return ids tokenized_sequences = [ _tok(prompt) for prompt in conversation_list ] if padding_side == "right": input_ids = torch.nn.utils.rnn.pad_sequence( tokenized_sequences, batch_first=True, padding_value=tokenizer.pad_token_id, ) else: tokenized_sequences = [seq.squeeze(0) for seq in tokenized_sequences] reversed_sequences = [seq.flip(0) for seq in tokenized_sequences] padded_reversed = torch.nn.utils.rnn.pad_sequence( reversed_sequences, batch_first=True, padding_value=tokenizer.pad_token_id ) input_ids = padded_reversed.flip(1) attention_masks = input_ids.ne(tokenizer.pad_token_id) # QFormer repeat logic expects exactly one question per conversation. if qformer_tokenizer is not None: if len(questions_list) != len(conversation_list): raise ValueError( "QFormer input misalignment: questions_list must match conversation_list " f"(got {len(questions_list)} questions vs {len(conversation_list)} conversations). " "Each dataset must return one question string per conversation." ) if len(questions_list) == 0: raise ValueError( "QFormer enabled but questions_list is empty. Each dataset must return questions." ) if qformer_tokenizer is not None: qformer_inputs = qformer_tokenizer( questions_list, padding=True, truncation=True, return_tensors="pt" ) qformer_input_ids = qformer_inputs.input_ids qformer_attention_masks = qformer_inputs.attention_mask else: qformer_input_ids = qformer_attention_masks = None # Preparing targets and handling conversation types conv = conversation_lib.default_conversation.copy() targets = input_ids.clone() # conv_type == "llava_v1" sep = conv.sep + conv.roles[1] + ": " sep2 = conv.sep2 for conversation, target in zip(conversation_list, targets): _process_conversation( conversation, target, tokenizer, sep, sep2, append_eos=append_eos ) # Adjusting for inferences if not inferences[0]: max_num_images = max((len(paths) for paths in image_path_list), default=0) truncate_len = tokenizer.model_max_length - (img_emb_len - 1) * max_num_images if input_ids.shape[1] > truncate_len: input_ids, targets, attention_masks = map( lambda x: x[:, :truncate_len], [input_ids, targets, attention_masks] ) # Guardrail: if a sample has GT masks but *no* [SEG] tokens survive truncation, the model will # produce no mask predictions and the loss path will assert. In that case, drop mask supervision # for the sample (set masks to None). If some [SEG] tokens survive, truncate GT masks to match. if seg_token_id is not None and len(masks_list) == len(batch): with torch.no_grad(): seg_counts_per_conv = ( (input_ids[:, 1:] == seg_token_id).sum(dim=1).to("cpu").tolist() if input_ids.numel() > 0 and input_ids.shape[1] > 1 else [0 for _ in range(input_ids.shape[0])] ) for sample_i in range(len(batch)): c0, c1 = int(conv_offset_list[sample_i]), int(conv_offset_list[sample_i + 1]) seg_count = int(sum(seg_counts_per_conv[c0:c1])) if c1 > c0 else 0 tensor_mask_idxs = [ j for j, m in enumerate(masks_list[sample_i]) if isinstance(m, torch.Tensor) and m.ndim >= 3 ] if not tensor_mask_idxs: continue if seg_count <= 0: for j in tensor_mask_idxs: masks_list[sample_i][j] = None continue # Only safe to truncate when there's a single tensor mask container (common single-image case). if len(tensor_mask_idxs) == 1: j = tensor_mask_idxs[0] m = masks_list[sample_i][j] if isinstance(m, torch.Tensor) and m.ndim >= 3 and m.shape[0] > seg_count: masks_list[sample_i][j] = m[:seg_count] if len(global_enc_image_list) > 0: global_enc_images_ret = torch.cat(global_enc_image_list, dim=0) else: global_enc_images_ret = None new_masks_list = list(zip_longest(*masks_list)) new_label_list = list(zip_longest(*label_list)) new_resize_list = list(zip_longest(*resize_list)) return { "image_paths": image_path_list, "global_enc_images": global_enc_images_ret, "grounding_enc_images": ( None if len(grounding_enc_image_list) == 0 or grounding_enc_image_list[0] is None else torch.cat(grounding_enc_image_list, dim=0) ), "input_ids": input_ids, "labels": targets, "attention_masks": attention_masks, "masks_list": new_masks_list, "label_list": new_label_list, "resize_list": new_resize_list, "conv_offset": torch.LongTensor(conv_offset_list), "image_offset": torch.LongTensor(img_offset_list), "grounding_offset": torch.LongTensor(grounding_offset_list), "questions_list": questions_list, "sampled_classes_list": selected_labels_list, "inference": inferences[0], "conversation_list": conversation_list, "qformer_input_ids": qformer_input_ids, "qformer_attention_masks": qformer_attention_masks, "sample_ids": sample_idxs, } def _process_conversation(conversation, target, tokenizer, sep, sep2, append_eos=True): total_len = target.ne(tokenizer.pad_token_id).sum().item() rounds = conversation.split(sep2) cur_len = 1 target[:cur_len] = IGNORE_INDEX def _len(text): if DEFAULT_IMAGE_TOKEN in conversation: ids = tokenizer_image_token(text, tokenizer) else: ids = tokenizer(text).input_ids if (not append_eos) and len(ids) > 0 and ids[-1] == tokenizer.eos_token_id: ids = ids[:-1] return len(ids) for rou in rounds: if not rou: break parts = rou.split(sep) if len(parts) == 1: round_len = _len(rou) - 1 target[cur_len : cur_len + round_len] = IGNORE_INDEX cur_len += round_len continue parts[0] += sep round_len = _len(rou) instruction_len = _len(parts[0]) - 2 target[cur_len : cur_len + instruction_len] = IGNORE_INDEX cur_len += round_len target[cur_len:] = IGNORE_INDEX if not append_eos and cur_len == total_len + 1: cur_len -= 1 if cur_len < tokenizer.model_max_length: assert cur_len == total_len, f"{cur_len = }, {total_len = }"