File size: 1,809 Bytes
cb634e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """Single-nucleotide tokenizer for genomic sequences.
Vocabulary (small, fixed):
0: <pad>
1: <bos>
2: <eos>
3: <mask>
4: A
5: C
6: G
7: T
8: N (unknown / soft-masked nucleotide)
Lower- and upper-case nucleotides collapse. Anything not in {A,C,G,T} maps to N.
"""
from __future__ import annotations
import numpy as np
import torch
PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
MASK_ID = 3
A_ID = 4
C_ID = 5
G_ID = 6
T_ID = 7
N_ID = 8
VOCAB_SIZE = 9
SPECIAL_IDS = (PAD_ID, BOS_ID, EOS_ID, MASK_ID)
NUCLEOTIDE_IDS = (A_ID, C_ID, G_ID, T_ID, N_ID)
def _build_lookup() -> np.ndarray:
"""Byte → token-id table. Length 256."""
table = np.full(256, N_ID, dtype=np.int64)
table[ord("A")] = A_ID
table[ord("a")] = A_ID
table[ord("C")] = C_ID
table[ord("c")] = C_ID
table[ord("G")] = G_ID
table[ord("g")] = G_ID
table[ord("T")] = T_ID
table[ord("t")] = T_ID
table[ord("U")] = T_ID
table[ord("u")] = T_ID
table[ord("N")] = N_ID
table[ord("n")] = N_ID
return table
_LOOKUP = _build_lookup()
def encode(seq: str | bytes) -> np.ndarray:
"""Encode a DNA string to a numpy int64 array of token-ids (no BOS/EOS)."""
if isinstance(seq, str):
seq = seq.encode("ascii", errors="replace")
arr = np.frombuffer(seq, dtype=np.uint8)
return _LOOKUP[arr]
def encode_torch(seq: str | bytes, device: str | torch.device = "cpu") -> torch.Tensor:
return torch.from_numpy(encode(seq).copy()).to(device)
_DECODE_TABLE = {
PAD_ID: "_",
BOS_ID: "[",
EOS_ID: "]",
MASK_ID: "?",
A_ID: "A",
C_ID: "C",
G_ID: "G",
T_ID: "T",
N_ID: "N",
}
def decode(ids) -> str:
if torch.is_tensor(ids):
ids = ids.detach().cpu().tolist()
return "".join(_DECODE_TABLE.get(int(i), "?") for i in ids)
|