File size: 6,006 Bytes
697dca3 | 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 112 113 114 115 116 117 | #!/usr/bin/env python3
"""HuggingFace slow-tokenizer wrapper for our morphological tokenizer, for the BabyLM leaderboard submission.
The official strict harness (`sentence_zero_shot/dataset.py`) calls `tokenizer(text, return_offsets_mapping=True)`
and locates the scored completion by CHARACTER offsets. A pure-Python tokenizer normally can't return offsets;
this subclass overrides the encode path to emit word-granularity `offset_mapping` (via OurTokenizer.encode_with_spans),
which is sufficient because completions always begin at a whitespace boundary.
Ship this file + our_vocab.json (+ the analyzer resources) with the model and load via `trust_remote_code=True`.
Special tokens are NOT auto-added (add_special_tokens is a no-op) to match the harness's no-BOS default for
from-scratch tokenizers.
"""
from __future__ import annotations
import json
from pathlib import Path
from transformers import PreTrainedTokenizer
from transformers.tokenization_utils_base import BatchEncoding
try:
# NOTE: HF's local-dir dynamic-module loader copies only THIS file's one level of relative
# imports (no recursion), so the analyzer must be imported here even though only
# our_tokenizer uses it.
from .morph_tokenizer import load_resources as _hf_copy_probe # noqa: F401
from .our_tokenizer import BOS, EOS, OurTokenizer, PAD, SPECIALS, UNK # noqa: F401
except ImportError: # direct use (importlib evades HF check_imports regex)
import importlib as _il
_m = _il.import_module("our_tokenizer")
BOS, EOS, OurTokenizer, PAD, SPECIALS, UNK = _m.BOS, _m.EOS, _m.OurTokenizer, _m.PAD, _m.SPECIALS, _m.UNK
class OurMorphTokenizer(PreTrainedTokenizer):
"""Slow HF tokenizer backed by OurTokenizer. Emits offset_mapping for the strict harness."""
vocab_files_names = {"vocab_file": "our_vocab.json"}
model_input_names = ["input_ids", "attention_mask"]
def __init__(self, vocab_file=None, char_backoff=True, **kwargs):
self._vocab = json.loads(Path(vocab_file).read_text(encoding="utf-8")) if vocab_file else {s: i for i, s in enumerate(SPECIALS)}
self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
self._ot = OurTokenizer.load(vocab_file, res=None) if vocab_file else None
if self._ot is not None:
self._ot.char_backoff = char_backoff
super().__init__(pad_token=PAD, unk_token=UNK, bos_token=BOS, eos_token=EOS,
char_backoff=char_backoff, **kwargs)
@property
def vocab_size(self) -> int:
return len(self._vocab)
def get_vocab(self) -> dict:
return dict(self._vocab)
def _tokenize(self, text, **kwargs):
return self._ot.tokenize(text)
def _convert_token_to_id(self, token):
return self._vocab.get(token, self._vocab[UNK])
def _convert_id_to_token(self, index):
return self._ids_to_tokens.get(index, UNK)
def convert_tokens_to_string(self, tokens):
return self._ot.decode([self._convert_token_to_id(t) for t in tokens])
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
# No auto special tokens (match the harness's no-BOS default); pairs are concatenated.
return token_ids_0 if token_ids_1 is None else token_ids_0 + token_ids_1
def save_vocabulary(self, save_directory, filename_prefix=None):
out = Path(save_directory) / ((filename_prefix + "-" if filename_prefix else "") + "our_vocab.json")
out.write_text(json.dumps(self._vocab, ensure_ascii=False), encoding="utf-8")
return (str(out),)
# -- offset-aware encoding (the reason this wrapper exists) -------------------------------------
def _one(self, text, return_offsets_mapping):
ids, spans = self._ot.encode_with_spans(text)
d = {"input_ids": ids, "attention_mask": [1] * len(ids)}
if return_offsets_mapping:
d["offset_mapping"] = spans
return d
def _encode_plus(self, text, text_pair=None, add_special_tokens=True, return_offsets_mapping=False,
return_attention_mask=None, return_tensors=None, **kwargs):
d = self._one(text, return_offsets_mapping)
if text_pair is not None:
d2 = self._one(text_pair, return_offsets_mapping)
d["input_ids"] += d2["input_ids"]
d["attention_mask"] += d2["attention_mask"]
if return_offsets_mapping:
d["offset_mapping"] += d2["offset_mapping"]
return BatchEncoding(d, tensor_type=return_tensors, prepend_batch_axis=True)
def _batch_encode_plus(self, batch_text_or_pairs, add_special_tokens=True, return_offsets_mapping=False,
return_attention_mask=None, return_tensors=None, padding_strategy=None, **kwargs):
rows = []
for t in batch_text_or_pairs:
text, pair = t if isinstance(t, tuple) else (t, None)
r = self._one(text, return_offsets_mapping)
if pair is not None:
r2 = self._one(pair, return_offsets_mapping)
r["input_ids"] += r2["input_ids"]; r["attention_mask"] += r2["attention_mask"]
if return_offsets_mapping:
r["offset_mapping"] += r2["offset_mapping"]
rows.append(r)
# left-as-lists (harness scores one at a time); pad only if a tensor type is requested
out = {k: [r[k] for r in rows] for k in rows[0]}
if return_tensors is not None:
mx = max(len(r["input_ids"]) for r in rows)
pad_id = self._vocab[PAD]
out["input_ids"] = [r["input_ids"] + [pad_id] * (mx - len(r["input_ids"])) for r in rows]
out["attention_mask"] = [r["attention_mask"] + [0] * (mx - len(r["attention_mask"])) for r in rows]
if return_offsets_mapping:
out["offset_mapping"] = [r["offset_mapping"] + [(0, 0)] * (mx - len(r["offset_mapping"])) for r in rows]
return BatchEncoding(out, tensor_type=return_tensors)
|