Promote Gene Tokens

Your task

Function signature

promote_gene_tokens(
    tokens: torch.Tensor,
    valid_mask: torch.Tensor,
    protected_mask: torch.Tensor,
    target_ids: torch.Tensor,
    target_mask: torch.Tensor,
    pad_token_id: int = 0,
) -> dict[str, torch.Tensor]

Worked example

Worked example

import torch

tokens = torch.tensor([[101, 11, 22, 33, 102, 0]], dtype=torch.int64)
valid_mask = torch.tensor([[True, True, True, True, True, False]], dtype=torch.bool)
protected_mask = torch.tensor([[True, False, False, False, True, False]], dtype=torch.bool)
target_ids = torch.tensor([[33, 44, 101]], dtype=torch.int64)
target_mask = torch.tensor([[True, True, True]], dtype=torch.bool)
pad_token_id = 0

result = promote_gene_tokens(tokens=tokens, valid_mask=valid_mask, protected_mask=protected_mask, target_ids=target_ids, target_mask=target_mask, pad_token_id=pad_token_id)

Expected return

{
    'tokens': torch.tensor([[101, 33, 44, 11, 22, 102]], dtype=torch.int64),
    'source_positions': torch.tensor([[0, 3, -1, 1, 2, 4]], dtype=torch.int64),
    'intervention_mask': torch.tensor([[False, True, True, False, False, False]], dtype=torch.bool),
    'valid_mask': torch.tensor([[True, True, True, True, True, True]], dtype=torch.bool),
    'lengths': torch.tensor([6], dtype=torch.int64),
    'target_found': torch.tensor([[True, False, True]], dtype=torch.bool),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• tokens is [B,L], with valid_mask forming a contiguous prefix. Valid IDs are unique, nonnegative, and unequal to pad_token_id (default 0). Active target IDs selected by target_mask obey the same uniqueness/nonpadding rules.
• protected_mask can mark only a boundary prefix and/or suffix of valid tokens; an all-protected sequence is treated as a protected prefix. Protected tokens cannot be edited, even when requested.
• Move requested mutable targets to the front of the mutable body IN REQUEST ORDER; insert requested IDs that are missing. Follow them with unrequested mutable tokens in original order. Each requested ID occurs once.
• Keep the protected prefix and suffix intact. If the edited body exceeds capacity, truncate only its tail; then append the suffix and pad the remaining slots.
• source_positions records the original index of each retained/moved token; inserted and padded slots use -1. intervention_mask is True only for retained promoted/inserted tokens. target_found reports active targets present ANYWHERE in the original valid sequence, including protected tokens.
• Return int64 token IDs, indices, and lengths, plus boolean masks. No floating differentiation is required; do not mutate the input sequence.

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 promote_gene_tokens(
    tokens: torch.Tensor,  # (B, L)
    valid_mask: torch.Tensor,  # (B, L)
    protected_mask: torch.Tensor,  # (B, L)
    target_ids: torch.Tensor,  # (B, K)
    target_mask: torch.Tensor,  # (B, K)
    pad_token_id: int = 0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    tokens: (B, L).
    source_positions: (B, L).
    intervention_mask: (B, L).
    valid_mask: (B, L).
    lengths: (B,).
    target_found: (B, K)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement promote_gene_tokens")

Related exercises

• Combined exerciseGene-token deletion and rank promotion

Sources and further reading

• www.nature.com
• huggingface.co
• huggingface.co