Delete Gene Tokens

Your task

Function signature

delete_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

Delete genes while keeping boundary tokens

import torch

tokens = torch.tensor([[91, 8, 4, 7, 92, 0], [5, 2, 9, 0, 0, 0]], dtype=torch.int64)
valid_mask = torch.tensor([[True, True, True, True, True, False], [True, True, True, False, False, False]], dtype=torch.bool)
protected_mask = torch.tensor([[True, False, False, False, True, False], [False, False, False, False, False, False]], dtype=torch.bool)
target_ids = torch.tensor([[4, 91, 30], [9, 6, -1]], dtype=torch.int64)
target_mask = torch.tensor([[True, True, True], [True, True, False]], dtype=torch.bool)
pad_token_id = 0

result = delete_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([[91, 8, 7, 92, 0, 0], [5, 2, 0, 0, 0, 0]], dtype=torch.int64),
    'source_positions': torch.tensor([[0, 1, 3, 4, -1, -1], [0, 1, -1, -1, -1, -1]], dtype=torch.int64),
    'intervention_mask': torch.full((2, 6), False, dtype=torch.bool),
    'valid_mask': torch.tensor([[True, True, True, True, False, False], [True, True, False, False, False, False]], dtype=torch.bool),
    'lengths': torch.tensor([4, 2], dtype=torch.int64),
    'target_found': torch.tensor([[True, True, False], [True, False, False]], 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.
• Remove active target IDs from the mutable body, keeping every other token in original order. Missing IDs have no effect.
• 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 (there are none in deletion mode). 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 delete_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 delete_gene_tokens")

Related exercises

• Combined exerciseGene-token deletion and rank promotion

Sources and further reading

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