Spaces:
Sleeping
Sleeping
File size: 3,163 Bytes
67f0f3c | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | from __future__ import annotations
from dataclasses import dataclass
from transformers import PreTrainedTokenizerBase
@dataclass(slots=True)
class TextChunk:
text: str
token_count: int
start_token: int
end_token: int
class TokenChunker:
"""Chunk text by model tokens while preserving original character spans."""
def __init__(
self,
tokenizer: PreTrainedTokenizerBase,
*,
max_input_tokens: int = 512,
overlap: int = 64,
text_prefix: str = "passage: ",
) -> None:
self.tokenizer = tokenizer
self.max_input_tokens = max_input_tokens
self.overlap = overlap
self.text_prefix = text_prefix
prefix_tokens = self.tokenizer(
text_prefix,
add_special_tokens=False,
return_attention_mask=False,
return_token_type_ids=False,
verbose=False,
)["input_ids"]
special_tokens = self.tokenizer.num_special_tokens_to_add(pair=False)
self.max_body_tokens = max(1, max_input_tokens - len(prefix_tokens) - special_tokens)
if overlap >= self.max_body_tokens:
raise ValueError(
f"overlap={overlap} must be smaller than usable body tokens={self.max_body_tokens}"
)
def count_tokens(self, text: str) -> int:
return len(
self.tokenizer(
text,
add_special_tokens=False,
return_attention_mask=False,
return_token_type_ids=False,
verbose=False,
)["input_ids"]
)
def split(self, text: str) -> list[TextChunk]:
text = text.strip()
if not text:
return []
encoded = self.tokenizer(
text,
add_special_tokens=False,
return_offsets_mapping=True,
return_attention_mask=False,
return_token_type_ids=False,
verbose=False,
)
input_ids = encoded["input_ids"]
offsets = encoded["offset_mapping"]
total_tokens = len(input_ids)
if total_tokens <= self.max_body_tokens:
return [
TextChunk(
text=text,
token_count=total_tokens,
start_token=0,
end_token=total_tokens,
)
]
chunks: list[TextChunk] = []
step = self.max_body_tokens - self.overlap
start = 0
while start < total_tokens:
end = min(start + self.max_body_tokens, total_tokens)
char_start = offsets[start][0]
char_end = offsets[end - 1][1]
chunk_text = text[char_start:char_end].strip()
if chunk_text:
chunks.append(
TextChunk(
text=chunk_text,
token_count=end - start,
start_token=start,
end_token=end,
)
)
if end >= total_tokens:
break
start += step
return chunks
|