Add truncation step to transformer model
@@ -6,7 +6,7 @@ from thinc.types import ArgsKwargs
from transformers.tokenization_utils import BatchEncoding
import logging
-from ..data_classes import FullTransformerBatch
+from ..data_classes import FullTransformerBatch, WordpieceBatch
from ..util import huggingface_tokenize, huggingface_from_pretrained
from ..util import find_last_hidden, maybe_flush_pytorch_cache
from ..truncate import truncate_oversize_splits
@@ -106,18 +106,19 @@ def forward(
maybe_flush_pytorch_cache(chance=model.attrs.get("flush_cache_chance", 0))
if "logger" in model.attrs:
log_gpu_memory(model.attrs["logger"], "begin forward")
- token_data = huggingface_tokenize(tokenizer, [span.text for span in flat_spans])
+ token_data = WordpieceBatch.from_batch_encoding(
+ huggingface_tokenize(tokenizer, [span.text for span in flat_spans])
+ )
if "logger" in model.attrs:
log_batch_size(model.attrs["logger"], token_data, is_train)
align = get_alignment(
flat_spans,
- token_data["input_texts"],
+ token_data.strings,
model.attrs["tokenizer"].all_special_tokens
)
token_data, align = truncate_oversize_splits(
token_data,
align,
- tokenizer.pad_token,
tokenizer.model_max_length
)
tensors, bp_tensors = transformer(token_data, is_train)
