File size: 3,574 Bytes
264f262 | 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 | """Extract complete DNA records from a PhaGen token canvas."""
from __future__ import annotations
from collections.abc import Sequence
from numbers import Integral
def extract_contigs(token_ids, tokenizer, min_length=1) -> list[str]:
"""Return ACGT strings from complete BOS/EOS spans in one token sequence.
``token_ids`` must be a one-dimensional integer sequence or an object such
as a one-dimensional torch tensor with a ``tolist()`` method. No torch
dependency is required. A nested BOS restarts the span; any non-ACGT token
inside a span discards it. Unclosed, empty and short spans are discarded.
``min_length`` counts DNA bases, excluding the boundary tokens.
Invalid input shapes, token types or length limits raise ``ValueError``.
The tokenizer must provide distinct BOS/EOS IDs and single-base ACGT IDs.
"""
if isinstance(min_length, bool) or not isinstance(min_length, Integral) or min_length < 1:
raise ValueError("min_length must be a positive integer counting DNA bases.")
if hasattr(token_ids, "ndim") and token_ids.ndim != 1:
raise ValueError("token_ids must be a single one-dimensional token sequence.")
if hasattr(token_ids, "tolist"):
try:
token_ids = token_ids.tolist()
except (TypeError, ValueError, RuntimeError) as exc:
raise ValueError("token_ids must convert to a one-dimensional integer sequence.") from exc
if isinstance(token_ids, (str, bytes, bytearray)) or not isinstance(token_ids, Sequence):
raise ValueError("token_ids must be a single one-dimensional integer sequence.")
if any(isinstance(token_id, bool) or not isinstance(token_id, Integral) for token_id in token_ids):
raise ValueError("token_ids must contain only integer token IDs, without nested sequences.")
bos_id = getattr(tokenizer, "bos_token_id", None)
eos_id = getattr(tokenizer, "eos_token_id", None)
if any(isinstance(token_id, bool) or not isinstance(token_id, Integral) for token_id in (bos_id, eos_id)):
raise ValueError("tokenizer must provide integer bos_token_id and eos_token_id values.")
if bos_id == eos_id:
raise ValueError("tokenizer BOS and EOS token IDs must be distinct.")
try:
base_ids = [tokenizer.convert_tokens_to_ids(base) for base in "ACGT"]
except (AttributeError, LookupError, TypeError, ValueError) as exc:
raise ValueError("tokenizer must map each of A, C, G and T to its own token ID.") from exc
special_ids = {
token_id
for name in ("bos_token_id", "eos_token_id", "pad_token_id", "mask_token_id", "unk_token_id")
if isinstance(token_id := getattr(tokenizer, name, None), Integral)
}
if (
any(isinstance(token_id, bool) or not isinstance(token_id, Integral) for token_id in base_ids)
or len(set(base_ids)) != 4
or set(base_ids) & special_ids
):
raise ValueError("tokenizer must provide four distinct ACGT token IDs separate from special tokens.")
base_by_id = dict(zip(base_ids, "ACGT"))
contigs = []
current = None
for token_id in token_ids:
if token_id == bos_id:
current = []
elif token_id == eos_id:
if current is not None and len(current) >= min_length:
contigs.append("".join(current))
current = None
elif current is not None:
if token_id in base_by_id:
current.append(base_by_id[token_id])
else:
current = None
return contigs
|