data/bio.py

from __future__ import annotations

PAD_ID = 256
CLS_ID = 257
SEP_ID = 258
BYTE_VOCAB = 259


def encode_with_offsets(text: str, max_len: int, add_cls: bool = True):
    raw = text.encode("utf-8", errors="replace")
    ids: list[int] = []
    offs: list[int] = []
    if add_cls:
        ids.append(CLS_ID)
        offs.append(-1)
    budget = max_len - len(ids)
    for j, b in enumerate(raw[:budget]):
        ids.append(int(b))
        offs.append(j)
    return (ids, offs, raw)


def char_span_to_byte_span(text: str, c_start: int, c_end: int) -> tuple[int, int]:
    b_start = len(text[:c_start].encode("utf-8"))
    b_end = len(text[:c_end].encode("utf-8"))
    return (b_start, b_end)


def spans_to_bio(text: str, spans, max_len: int, tag_to_id: dict) -> tuple[list, list]:
    ids, offs, _ = encode_with_offsets(text, max_len)
    o_id = tag_to_id["O"]
    tags = [o_id] * len(ids)
    for b_start, b_end, entity in sorted(spans, key=lambda s: s[0]):
        b_tag = tag_to_id.get(f"B-{entity}")
        i_tag = tag_to_id.get(f"I-{entity}")
        if b_tag is None or i_tag is None:
            continue
        first = True
        for pos, off in enumerate(offs):
            if off < 0:
                continue
            if b_start <= off < b_end:
                tags[pos] = b_tag if first else i_tag
                first = False
    return (ids, tags)