Iterative masked-token reveal with fixed prompts

Your task

Function signature

iterative_masked_token_update(
    tokens: torch.Tensor,
    proposals: torch.Tensor,
    probabilities: torch.Tensor,
    priorities: torch.Tensor,
    lengths: torch.Tensor,
    position_mask: torch.Tensor,
    editable_mask: torch.Tensor,
    target_remaining: torch.Tensor,
    strategy: str = 'entropy',
    mask_id: int = -1,
) -> dict[str, torch.Tensor]

Worked example

Reveal low entropy positions until one mask remains

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• tokens and proposals are integer [B,L,D], with B,D\ge 1,L\ge 2; D=1 represents a sequence/structure track and D=8 can represent function-code slots. probabilities [B,L,D,V] has V\ge 2, finite NONNEGATIVE supplied categorical probabilities, each slot summing to 1.
• This independently supplied confidence distribution alone determines entropy; it need not be the categorical law that produced proposals. No universal upstream before/after-top-p distribution is assumed. Zero probabilities may represent excluded vocabulary IDs. All floats share dtype/device.
• priorities [B,L] is supplied finite random priority, not drawn internally. lengths [B] integers lie in [2,L] and include BOS/EOS: position 0 and position \mathrm{lengths}[b]-1 are special, and all later positions are padding.
• position_mask [B,L] and editable_mask [B,L] are bool with True means present/eligible-for-editing respectively. target_remaining [B] is a nonnegative integer count supplied by the schedule. Config mask_id is an integer (default -1); strategy is 'entropy' (default) or 'random'.
• A candidate position must satisfy 0<i<\mathrm{length}-1, position_mask and editable_mask True, and ALL D current slots equal mask_id. Thus partially known function slots are not overwritten.
• For every valid nonspecial position return \mathrm{entropy}_i=-\sum_{\mathrm{slot},v}p_{i,\mathrm{slot},v}\log p_{i,\mathrm{slot},v}, defining a p=0 term as 0 with a chosen zero derivative on that branch; its singular ambient derivative at zero is not claimed. Invalid/special/padded entropy is exactly 0.
• For strategy entropy, sort candidate positions by increasing total entropy; for random, by increasing supplied priorities. Equal scores are ordered by smaller position index. The number to commit is \operatorname{max}(\mathrm{number\_of\_candidates}-\mathrm{target\_remaining},0), so an oversized remaining target safely commits none.
• Replace ALL D slots at committed positions with their supplied proposals, preserving every other integer exactly. Proposal IDs are in [0,V) and must differ from mask_id at committed positions, ensuring a monotonic reveal.
• Return tokens, commit_mask [B,L], remaining_mask [B,L] for eligible positions still masked, committed_count [B], remaining_count [B], and entropy [B,L]. No neural-model call, categorical draw, remasking of known tokens, or schedule evaluation is hidden in this step.
• Batch rows have independent lengths and budgets; no-mask and all-fixed cases are allowed.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def iterative_masked_token_update(
    tokens: torch.Tensor,  # (B, L, D)
    proposals: torch.Tensor,  # (B, L, D)
    probabilities: torch.Tensor,  # (B, L, D, V)
    priorities: torch.Tensor,  # (B, L)
    lengths: torch.Tensor,  # (B,)
    position_mask: torch.Tensor,  # (B, L)
    editable_mask: torch.Tensor,  # (B, L)
    target_remaining: torch.Tensor,  # (B,)
    strategy: str = 'entropy',
    mask_id: int = -1,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    tokens: (B, L, D).
    commit_mask: (B, L).
    remaining_mask: (B, L).
    committed_count: (B,).
    remaining_count: (B,).
    entropy: (B, L)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement iterative_masked_token_update")

Sources and further reading

• www.science.org
• Biohub/esm · generation.py
• Biohub/esm · generation.py
• Biohub/esm · generation.py