Spaces:
Running
Running
File size: 1,600 Bytes
6d86da0 | 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 | """Tokenizer-dependent helpers shared by both stages."""
from __future__ import annotations
from dataclasses import dataclass
from .bio import content_offsets
from .spans import Span
def special_ids(tokenizer) -> tuple[int, int, int]:
cls = tokenizer.cls_token_id if tokenizer.cls_token_id is not None else tokenizer.bos_token_id
sep = tokenizer.sep_token_id if tokenizer.sep_token_id is not None else tokenizer.eos_token_id
pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
if cls is None or sep is None:
raise ValueError("Tokenizer needs CLS/SEP (or BOS/EOS) tokens")
return cls, sep, pad
def require_fast(tokenizer) -> None:
if not getattr(tokenizer, "is_fast", False):
raise RuntimeError("A fast tokenizer is required for offset mapping "
"(pip install sentencepiece protobuf tokenizers)")
@dataclass
class ArticleEncoding:
article_id: str
ids: list[int]
offsets: list[Span] # whitespace-stripped char offsets, one per token
def encode_article(tokenizer, article_id: str, text: str) -> ArticleEncoding:
"""Tokenize the whole article once, without special tokens or truncation.
Windows are cut from this later, so token indices are global and window
predictions can be merged unambiguously.
"""
enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True,
truncation=False, verbose=False)
return ArticleEncoding(article_id, list(enc["input_ids"]),
content_offsets(text, enc["offset_mapping"]))
|