gap_sdar_training.py

from dataclasses import dataclass
from typing import Dict, Optional, Sequence, Union

import torch
import torch.nn.functional as F


@dataclass
class GapRemaskOutputs:
    full_candidate_mask: torch.BoolTensor
    remask_target_flat: torch.BoolTensor
    remask_pred_full: torch.BoolTensor
    z_accept: torch.LongTensor
    z_proj: torch.LongTensor
    projected_mask: torch.BoolTensor
    projected_p_mask: torch.FloatTensor
    remask_loss: torch.Tensor
    metrics: Dict[str, float]


def _scatter_flat_mask(base_mask: torch.BoolTensor, selected_flat_mask: torch.BoolTensor) -> torch.BoolTensor:
    full_mask = torch.zeros_like(base_mask)
    full_mask[base_mask] = selected_flat_mask
    return full_mask


def build_p_mask_full(
    masked_indices: torch.BoolTensor,
    p_mask: torch.FloatTensor,
    shape: torch.Size,
    default_p_mask: float,
) -> torch.FloatTensor:
    p_mask_full = torch.full(shape, default_p_mask, dtype=torch.float32, device=masked_indices.device)
    p_mask_full[masked_indices] = p_mask.float()
    return p_mask_full


def get_num_transfer_tokens(block_length: int, steps: int) -> torch.LongTensor:
    if steps <= 0:
        raise ValueError(f"steps must be positive, got {steps}")

    base = block_length // steps
    remainder = block_length % steps
    num_transfer_tokens = torch.full((steps,), base, dtype=torch.long)
    num_transfer_tokens[:remainder] += 1
    return num_transfer_tokens


def _select_block_positions(
    block_scores: torch.FloatTensor,
    masked_local_indices: torch.LongTensor,
    num_transfer_tokens: int,
    strategy: str,
    confidence_threshold: float,
) -> torch.LongTensor:
    k = min(num_transfer_tokens, int(masked_local_indices.numel()))
    if k <= 0:
        return masked_local_indices[:0]

    if strategy == "low_confidence_dynamic":
        high_conf_mask = block_scores > confidence_threshold
        if int(high_conf_mask.sum().item()) >= num_transfer_tokens:
            return masked_local_indices[high_conf_mask]

        topk = torch.topk(block_scores, k=k, sorted=False).indices
        return masked_local_indices[topk]
    if strategy == "low_confidence_static":
        topk = torch.topk(block_scores, k=k, sorted=False).indices
        return masked_local_indices[topk]
    if strategy == "sequential":
        return masked_local_indices[:k]

    raise ValueError(f"Unsupported rollout strategy: {strategy}")


def _resolve_num_transfer_tokens(
    num_transfer_tokens: Union[int, torch.Tensor, Sequence[int]],
    batch_idx: int,
) -> int:
    if torch.is_tensor(num_transfer_tokens):
        if num_transfer_tokens.numel() == 1:
            return int(num_transfer_tokens.item())
        return int(num_transfer_tokens[batch_idx].item())

    if isinstance(num_transfer_tokens, Sequence) and not isinstance(num_transfer_tokens, (str, bytes)):
        if len(num_transfer_tokens) == 1:
            return int(num_transfer_tokens[0])
        return int(num_transfer_tokens[batch_idx])

    return int(num_transfer_tokens)


def select_policy_transfer_tokens(
    masked_indices: torch.BoolTensor,
    proposal_scores_full: torch.FloatTensor,
    num_tokens,
    block_size: int,
    num_transfer_tokens: Union[int, torch.Tensor, Sequence[int]],
    strategy: str = "low_confidence_dynamic",
    confidence_threshold: float = 0.95,
    scope: str = "all",
) -> torch.BoolTensor:
    if scope not in {"all", "frontier_block"}:
        raise ValueError(f"Unsupported rollout scope: {scope}")

    reveal_mask = torch.zeros_like(masked_indices)
    if torch.is_tensor(num_transfer_tokens):
        if int(num_transfer_tokens.max().item()) <= 0:
            return reveal_mask
    elif isinstance(num_transfer_tokens, Sequence) and not isinstance(num_transfer_tokens, (str, bytes)):
        if max(int(x) for x in num_transfer_tokens) <= 0:
            return reveal_mask
    elif int(num_transfer_tokens) <= 0:
        return reveal_mask

    for batch_idx, packed_lengths in enumerate(num_tokens):
        current_num_transfer_tokens = _resolve_num_transfer_tokens(num_transfer_tokens, batch_idx)
        cursor = 0
        for sample_len_tensor in packed_lengths:
            sample_len = int(sample_len_tensor.item())
            sample_end = cursor + sample_len

            for block_start in range(cursor, sample_end, block_size):
                block_end = min(block_start + block_size, sample_end)
                block_mask = masked_indices[batch_idx, block_start:block_end]
                if not block_mask.any():
                    continue

                masked_local_indices = torch.nonzero(block_mask, as_tuple=False).flatten()
                block_scores = proposal_scores_full[batch_idx, block_start:block_end][masked_local_indices]
                chosen = _select_block_positions(
                    block_scores=block_scores,
                    masked_local_indices=masked_local_indices,
                    num_transfer_tokens=current_num_transfer_tokens,
                    strategy=strategy,
                    confidence_threshold=confidence_threshold,
                )
                reveal_mask[batch_idx, block_start:block_end][chosen] = True

                if scope == "frontier_block":
                    break

            cursor = sample_end

    return reveal_mask


def select_teacher_forced_rollout_tokens(
    masked_indices: torch.BoolTensor,
    proposal_scores_full: torch.FloatTensor,
    num_tokens,
    block_size: int,
    num_transfer_tokens: Union[int, torch.Tensor, Sequence[int]],
    strategy: str = "low_confidence_dynamic",
    confidence_threshold: float = 0.95,
    scope: str = "all",
) -> torch.BoolTensor:
    return select_policy_transfer_tokens(
        masked_indices=masked_indices,
        proposal_scores_full=proposal_scores_full,
        num_tokens=num_tokens,
        block_size=block_size,
        num_transfer_tokens=num_transfer_tokens,
        strategy=strategy,
        confidence_threshold=confidence_threshold,
        scope=scope,
    )


def build_rollout_scope_mask(
    masked_indices: torch.BoolTensor,
    reference_mask: torch.BoolTensor,
    num_tokens,
    block_size: int,
    scope: str = "all",
) -> torch.BoolTensor:
    if scope not in {"all", "frontier_block"}:
        raise ValueError(f"Unsupported rollout scope: {scope}")

    if scope == "all":
        return reference_mask.clone()

    scope_mask = torch.zeros_like(reference_mask)
    for batch_idx, packed_lengths in enumerate(num_tokens):
        cursor = 0
        for sample_len_tensor in packed_lengths:
            sample_len = int(sample_len_tensor.item())
            sample_end = cursor + sample_len

            for block_start in range(cursor, sample_end, block_size):
                block_end = min(block_start + block_size, sample_end)
                if not masked_indices[batch_idx, block_start:block_end].any():
                    continue

                scope_mask[batch_idx, block_start:block_end] = reference_mask[batch_idx, block_start:block_end]
                break

            cursor = sample_end

    return scope_mask


def build_rollout_p_mask(
    masked_indices: torch.BoolTensor,
    labels: torch.LongTensor,
    num_tokens,
    target_scope_mask: Optional[torch.BoolTensor] = None,
    per_block: bool = False,
    block_size: Optional[int] = None,
    eps: float = 1e-3,
) -> torch.FloatTensor:
    p_mask_full = torch.full(masked_indices.shape, eps, dtype=torch.float32, device=masked_indices.device)

    for batch_idx, packed_lengths in enumerate(num_tokens):
        cursor = 0
        for sample_len_tensor in packed_lengths:
            sample_len = int(sample_len_tensor.item())
            sample_end = cursor + sample_len
            if per_block:
                if block_size is None:
                    raise ValueError("block_size must be provided when per_block=True")
                for block_start in range(cursor, sample_end, block_size):
                    block_end = min(block_start + block_size, sample_end)
                    block_target_mask = labels[batch_idx, block_start:block_end].ne(-100)
                    if target_scope_mask is not None:
                        block_target_mask = block_target_mask & target_scope_mask[batch_idx, block_start:block_end]
                    target_count = int(block_target_mask.sum().item())
                    if target_count == 0:
                        continue
                    block_masked = masked_indices[batch_idx, block_start:block_end] & block_target_mask
                    block_p_mask = max(block_masked.sum().item() / target_count, eps)
                    p_mask_full[batch_idx, block_start:block_end][block_masked] = block_p_mask
            else:
                sample_target_mask = labels[batch_idx, cursor:sample_end].ne(-100)
                if target_scope_mask is not None:
                    sample_target_mask = sample_target_mask & target_scope_mask[batch_idx, cursor:sample_end]
                target_count = int(sample_target_mask.sum().item())
                if target_count > 0:
                    sample_masked = masked_indices[batch_idx, cursor:sample_end] & sample_target_mask
                    sample_p_mask = max(sample_masked.sum().item() / target_count, eps)
                    p_mask_full[batch_idx, cursor:sample_end][sample_masked] = sample_p_mask
            cursor = sample_end

    return p_mask_full[masked_indices]


def _expand_positive_blocks(
    positive_mask: torch.BoolTensor,
    candidate_mask: torch.BoolTensor,
    block_size: int,
) -> torch.BoolTensor:
    expanded = positive_mask.clone()
    _, seq_len = positive_mask.shape
    for block_start in range(0, seq_len, block_size):
        block_end = min(block_start + block_size, seq_len)
        block_positive = positive_mask[:, block_start:block_end].any(dim=1, keepdim=True)
        if not bool(block_positive.any().item()):
            continue
        expanded[:, block_start:block_end] |= candidate_mask[:, block_start:block_end] & block_positive
    return expanded


def _build_remask_targets(
    masked_indices: torch.BoolTensor,
    proposal_ids: torch.LongTensor,
    clean_targets_flat: torch.LongTensor,
    full_candidate_mask: torch.BoolTensor,
    block_size: int,
    supervision: Optional[str],
) -> torch.BoolTensor:
    supervision = (supervision or "adv_bce").strip().lower()
    wrong_flat_all = proposal_ids.ne(clean_targets_flat)
    wrong_mask = _scatter_flat_mask(masked_indices, wrong_flat_all) & full_candidate_mask
    if supervision in {"gt_mismatch_block_bce", "block_mismatch_bce", "gt_block_bce"}:
        return _expand_positive_blocks(wrong_mask, full_candidate_mask, block_size)
    return wrong_mask


def apply_gap_remask(
    noisy_input_ids: torch.LongTensor,
    clean_input_ids: torch.LongTensor,
    labels: torch.LongTensor,
    masked_indices: torch.BoolTensor,
    p_mask: torch.FloatTensor,
    proposal_ids: torch.LongTensor,
    remask_logits: torch.FloatTensor,
    candidate_mask_full: torch.BoolTensor,
    mask_token_id: int,
    remask_threshold: float,
    remask_loss_weight: float,
    remask_default_p_mask: float,
    block_size: int,
    supervision: Optional[str] = None,
    target_scope_mask: Optional[torch.BoolTensor] = None,
    ignore_index: int = -100,
) -> GapRemaskOutputs:
    full_candidate_mask = candidate_mask_full & masked_indices
    candidate_mask_flat = full_candidate_mask[masked_indices]

    clean_targets_flat = clean_input_ids[masked_indices]
    remask_target_full = _build_remask_targets(
        masked_indices=masked_indices,
        proposal_ids=proposal_ids,
        clean_targets_flat=clean_targets_flat,
        full_candidate_mask=full_candidate_mask,
        block_size=block_size,
        supervision=supervision,
    )
    remask_target_flat = remask_target_full[masked_indices]

    z_accept = noisy_input_ids.clone()
    if candidate_mask_flat.any():
        z_accept[full_candidate_mask] = clean_input_ids[full_candidate_mask]

    candidate_logits = remask_logits[candidate_mask_flat]
    candidate_targets = remask_target_flat[candidate_mask_flat].float()
    remask_pred_flat = torch.zeros_like(candidate_mask_flat)
    pos_weight_value = 1.0
    if candidate_logits.numel() > 0:
        positive_count = float(candidate_targets.sum().item())
        negative_count = float(candidate_targets.numel() - positive_count)
        pos_weight = None
        if positive_count > 0.0 and negative_count > 0.0:
            pos_weight_value = max(1.0, min(8.0, negative_count / positive_count))
            pos_weight = candidate_logits.new_tensor(pos_weight_value)
        remask_loss = F.binary_cross_entropy_with_logits(
            candidate_logits,
            candidate_targets,
            pos_weight=pos_weight,
        )
        remask_pred_flat[candidate_mask_flat] = torch.sigmoid(candidate_logits) >= remask_threshold
    else:
        remask_loss = remask_logits.sum() * 0.0

    remask_pred_full = _scatter_flat_mask(masked_indices, remask_pred_flat)
    z_proj = z_accept.clone()
    z_proj[remask_pred_full] = mask_token_id

    if target_scope_mask is None:
        target_scope_mask = labels.ne(ignore_index)
    else:
        target_scope_mask = target_scope_mask & labels.ne(ignore_index)

    projected_mask = z_proj.eq(mask_token_id) & target_scope_mask
    if not projected_mask.any():
        fallback_mask = full_candidate_mask & target_scope_mask
        if not fallback_mask.any():
            fallback_mask = masked_indices & target_scope_mask
        if not fallback_mask.any():
            fallback_mask = target_scope_mask
        fallback_indices = torch.nonzero(fallback_mask, as_tuple=False)
        if fallback_indices.numel() > 0:
            row, col = fallback_indices[0].tolist()
            z_proj[row, col] = mask_token_id
            projected_mask[row, col] = True

    p_mask_full = build_p_mask_full(
        masked_indices=masked_indices,
        p_mask=p_mask,
        shape=noisy_input_ids.shape,
        default_p_mask=remask_default_p_mask,
    )
    projected_p_mask = p_mask_full[projected_mask]

    candidate_total = max(int(candidate_mask_flat.sum().item()), 1)
    remask_positive = int(remask_target_flat.sum().item())
    remask_predicted = int(remask_pred_full.sum().item())
    remask_true_positive = int((remask_pred_full & remask_target_full).sum().item())
    precision = remask_true_positive / max(remask_predicted, 1)
    recall = remask_true_positive / max(remask_positive, 1)
    metrics = {
        "candidate_tokens": float(candidate_mask_flat.sum().item()),
        "remask_positive_rate": remask_positive / candidate_total,
        "remask_pred_rate": remask_predicted / candidate_total,
        "remask_precision": precision,
        "remask_recall": recall,
        "remask_pos_weight": pos_weight_value,
        "projected_mask_tokens": float(projected_mask.sum().item()),
        "remask_loss": float((remask_loss.detach() * remask_loss_weight).item()),
    }

    return GapRemaskOutputs(
        full_candidate_mask=full_candidate_mask,
        remask_target_flat=remask_target_flat,
        remask_pred_full=remask_pred_full,
        z_accept=z_accept,
        z_proj=z_proj,
        projected_mask=projected_mask,
        projected_p_mask=projected_p_mask,
        remask_loss=remask_loss * remask_loss_weight,
        metrics=metrics,
    )