Keep the PRIMA architecture and verify prediction resume provenance

06b25e4 verified about 9 hours ago

12.7 kB

"""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 = }"

	"""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 = }"