Problem

Three small accuracy issues in src/alpamayo_r1/utils/get_label_mask.py:

1. Truncated docstring in fill_masks_between_special_tokens — the entry for labels_mask cuts off mid-sentence:

   labels_mask (B, seq_len): tensor of labels mask, indicating the positions that should
   

   Readers can't tell what the function actually does to the mask, and the in-place mutation isn't documented.

2. Wrong return-type annotation on get_assistant_mask. The function declares -> torch.Tensor, but the body returns masks.tolist() (a list[bool]) when tokens is a list[int]:

   if isinstance(tokens, torch.Tensor):
       return torch.from_numpy(masks)
   else:
       return masks.tolist()

   This trips up type-checkers and is misleading to callers passing a list.

3. Dangling reference block in get_assistant_mask — the docstring ends with:

   Reference:
       Adapted from:
   

   No URL or citation follows. Either fill it in or drop it; this PR drops it.

Fix

Pure docstring + annotation hygiene. No behavioral change.

• Complete the truncated labels_mask description and note in-place mutation + inclusive end-token semantics (matches the # NOTE: we should include the end token in the mask comment that already exists in the body).
• Update get_assistant_mask annotation to torch.Tensor | list[bool] and reflect this in the docstring.
• Remove the empty Reference: Adapted from: block.

Verification

No code paths changed. Mypy / pyright now agree with the actual runtime types.