NeTSlab/gpt2_parfind_en_zh_equal
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: NeTSlab/gpt2_parfind_en_zh_equal
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-07-18 01:20:11 +08:00
commit 1ae14a5e92
17 changed files with 84476 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
.gitattributes vendored Normal file
View File

@@ -0,0 +1,35 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text

30
README.md Normal file
View File

@@ -0,0 +1,30 @@
---
language:
- en
- zh
tags:
- causal-lm
- language-model
- babylm
- babylm-2026
- multilingual
- gpt2
- paradigmfinder
pipeline_tag: text-generation
library_name: transformers
---
# gpt2_parfind_en_zh_equal
Bilingual GPT-2 model packaged for the BabyLM 2026 multilingual evaluation track.
Source experiment: `/home/achille.fusco/pr_baby_lm/BabyLM_2026_ENH/04-experiments/model_gpt2_ParFindFast_eng_zho_BD_budget16k_zhchildes_v1.0`
Hugging Face target repo: `NeTSlab/gpt2_parfind_en_zh_equal`
Current status on July 6, 2026: `resume_running`
Notes:
- `main` is intended to point to the final 1000M checkpoint (`epoch_9`).
- The custom ParadigmFinder tokenizer is bundled with the model files.
- Loading from HF requires `trust_remote_code=True`.

1
__init__.py Normal file
View File

@@ -0,0 +1 @@
__all__ = []

49
boundary_discovery.py Normal file
View File

@@ -0,0 +1,49 @@
import re
from typing import Iterable, List
DEFAULT_ANCHOR_RE = re.compile(r"[.!?;:,…·。！？；：，、（）()\[\]{}\"'«»\n\r\t]+")
def anchor_sequences(
text: str,
space_marker: str = "_",
min_sequence_length: int = 2,
) -> List[str]:
"""Return true-anchored sequences with internal spaces preserved as markers.
This mirrors the MorPiece boundary-discovery preparation step closely:
split only on strong punctuation/newline anchors, keep spaces inside a span
as a soft cue rather than a delimiter, and drop very short fragments.
"""
pieces = []
for seg in DEFAULT_ANCHOR_RE.split(text):
if not seg:
continue
filtered = "".join(ch for ch in seg if ch.isalpha() or ch in (" ", "'", "-"))
filtered = re.sub(r"\s+", " ", filtered).strip()
if len(filtered) < min_sequence_length:
continue
pieces.append(filtered.replace(" ", space_marker))
return pieces
def collect_boundary_units(
lines: Iterable[str],
space_marker: str = "_",
min_sequence_length: int = 2,
shorter_first: bool = True,
) -> List[str]:
"""Collect anchored training units from a raw-text iterator."""
units = []
for line in lines:
units.extend(
anchor_sequences(
line,
space_marker=space_marker,
min_sequence_length=min_sequence_length,
)
)
if shorter_first:
units.sort(key=len)
return units

32
config.json Normal file
View File

@@ -0,0 +1,32 @@
{
"activation_function": "gelu_new",
"architectures": [
"GPT2LMHeadModel"
],
"attn_pdrop": 0.1,
"bos_token_id": 2,
"embd_pdrop": 0.1,
"eos_token_id": 3,
"initializer_range": 0.02,
"layer_norm_epsilon": 1e-05,
"model_type": "gpt2",
"n_embd": 768,
"n_head": 12,
"n_inner": null,
"n_layer": 12,
"n_positions": 1024,
"pad_token_id": 0,
"reorder_and_upcast_attn": false,
"resid_pdrop": 0.1,
"scale_attn_by_inverse_layer_idx": false,
"scale_attn_weights": true,
"summary_activation": null,
"summary_first_dropout": 0.1,
"summary_proj_to_labels": true,
"summary_type": "cls_index",
"summary_use_proj": true,
"torch_dtype": "float32",
"transformers_version": "4.53.2",
"use_cache": true,
"vocab_size": 32219
}

7
generation_config.json Normal file
View File

@@ -0,0 +1,7 @@
{
"_from_model_config": true,
"bos_token_id": 2,
"eos_token_id": 3,
"pad_token_id": 0,
"transformers_version": "4.53.2"
}

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:b7c5bc6c3e5bc88bd7695aaa346a9acdea08497b1601d8936d70e1911e188019
size 442361472

20
multilingual_meta.json Normal file
View File

@@ -0,0 +1,20 @@
{
"mode": "soft",
"budget_per_language": 16384,
"languages": [
"eng",
"zho"
],
"tokens_per_language": {
"eng": 16384,
"zho": 16231
},
"shared_token_instances_deduped": 93,
"final_vocab_size": 32526,
"space_free_lexicon_languages": [
"zho"
],
"space_free_lexicon_token_counts": {
"zho": 16231
}
}

430
paradigm_utils.py Normal file
View File

@@ -0,0 +1,430 @@
# paradigm_utils.py
import time
from collections import defaultdict
from tqdm import tqdm
import os
import math
import json
from typing import List, Tuple, Set, Dict, Any
def _serialize_suffixes(sfx_set):
flat = []
for s in sfx_set:
if isinstance(s, tuple):
base, nested = s
flat.append([base, sorted(list(nested))]) # JSON-safe pair
else:
flat.append(s) # plain string
# stable order: strings first, then pairs; then lexicographic
def key(x):
return (0, x) if isinstance(x, str) else (1, x[0], tuple(x[1]))
return sorted(flat, key=key)
def paradigms_to_json(paradigms):
out = []
for stems, suffixes in paradigms:
out.append({
"stems": sorted(list(stems)),
"suffixes": _serialize_suffixes(suffixes),
})
return out
def save_paradigms_json(paradigms, path, meta=None):
payload = {
"schema_version": 1,
"created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
"meta": meta or {},
"paradigms": paradigms_to_json(paradigms),
}
with open(path, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
def _deserialize_suffixes(sfx_list):
out = set()
for item in sfx_list:
if isinstance(item, list): # [base, nested_list]
base, nested = item
out.add((base, frozenset(nested)))
else:
out.add(item)
return out
def load_paradigms_json(path):
with open(path, "r", encoding="utf-8") as f:
payload = json.load(f)
paradigms = []
for p in payload["paradigms"]:
stems = set(p["stems"])
suffixes = _deserialize_suffixes(p["suffixes"])
paradigms.append((stems, suffixes))
meta = payload.get("meta", {})
return paradigms, meta
### -----------------------------
### 1. Extract (stem, suffix) pairs from vocabulary
### -----------------------------
def extract_stem_suffix_pairs(vocab):
"""Return a mapping from stems to all suffixes they occur with, including null suffix."""
stem_to_suffixes = defaultdict(set)
for word in tqdm(vocab, desc="[1/7] Extracting stem-suffix pairs"):
for i in range(0, len(word) + 1): # include empty suffix
stem, suffix = word[:i], word[i:]
stem_to_suffixes[stem].add(suffix)
return stem_to_suffixes
### -----------------------------
### 2. Group stems by shared suffix sets and normalize by common prefix
### -----------------------------
def group_stems_by_suffixes(stem_to_suffixes, min_shared_stems=2, min_suffixes=2):
suffix_to_stems = defaultdict(set)
for stem, suffixes in stem_to_suffixes.items():
suffix_key = frozenset(suffixes)
suffix_to_stems[suffix_key].add(stem)
normalized_suffix_map = defaultdict(set)
for suffixes, stems in tqdm(suffix_to_stems.items(), desc="[2/7] Grouping and normalizing"):
non_empty_suffixes = [s for s in suffixes if s]
if len(stems) >= min_shared_stems and len(suffixes) >= min_suffixes:
common_prefix = os.path.commonprefix(non_empty_suffixes) if non_empty_suffixes else ""
if common_prefix:
normalized_stems = {stem + common_prefix for stem in stems}
adjusted_suffixes = {s[len(common_prefix):] if s.startswith(common_prefix) else s for s in suffixes}
else:
normalized_stems = stems
adjusted_suffixes = suffixes
if len(adjusted_suffixes) >= min_suffixes:
suffix_key = frozenset(adjusted_suffixes)
normalized_suffix_map[suffix_key].update(normalized_stems)
paradigms = [(stems, set(suffixes)) for suffixes, stems in normalized_suffix_map.items()]
return paradigms
### -----------------------------
### 3. Expand stem sets based on suffix set coverage
### -----------------------------
def stem_set_expansion(paradigms, stem_to_suffixes):
updated = 0
suffix_to_stems = {frozenset(suffixes): set(stems) for stems, suffixes in paradigms}
for stem, suffixes in tqdm(stem_to_suffixes.items(), desc="[3/7] Expanding stem sets"):
added = False
for paradigm_suffixes in sorted(suffix_to_stems.keys(), key=lambda x: (-len(x), tuple(sorted(x)))):
if paradigm_suffixes.issubset(suffixes):
if stem not in suffix_to_stems[paradigm_suffixes]:
suffix_to_stems[paradigm_suffixes].add(stem)
updated += 1
added = True
if not added and stem == 'design':
print(f"[DEBUG] No suitable paradigm for 'design' with suffixes {suffixes}")
enriched = [(stems, set(suffixes)) for suffixes, stems in suffix_to_stems.items()]
print(f"✅ Added {updated} stems via stem set expansion.")
return enriched
### -----------------------------
### 4. Expand suffix sets based on partial compatibility
### -----------------------------
def harmonic_number(n):
return sum(1.0 / i for i in range(1, n + 1))
def suffix_set_expansion(paradigms):
base = paradigms[:] # snapshot
merged = [ (set(stems), set(suffixes)) for stems, suffixes in base ]
enriched_count = 0
# Iterate in a deterministic order
for i, (stems_i, suffixes_i) in enumerate(sort_paradigms(merged)):
for j, (stems_j, suffixes_j) in enumerate(sort_paradigms(merged)):
if i == j:
continue
if suffixes_i > suffixes_j:
intersection = stems_i & stems_j
denom = max(1, len(stems_j)) # guard
if (len(stems_j) - len(intersection)) < (len(stems_j) / harmonic_number(denom)):
stems_i |= stems_j
enriched_count += 1
# do not mutate stems_j/suffixes_j further
print(f"\n✅ Enriched {enriched_count} paradigms via suffix set expansion.")
# Return back in original tuple-of-sets form
return [ (set(st), set(sf)) for st, sf in sort_paradigms(merged) ]
### -----------------------------
### 5. Prune subsumed stems
### -----------------------------
def prune_subsumed_stems(paradigms):
pruned_paradigms = []
for i, (stems_i, suffixes_i) in enumerate(paradigms):
pruned_stems = set(stems_i)
for j, (stems_j, suffixes_j) in enumerate(paradigms):
if i == j:
continue
if suffixes_j >= suffixes_i:
pruned_stems -= (stems_j & stems_i)
if pruned_stems:
pruned_paradigms.append((pruned_stems, suffixes_i))
print(f"✅ Pruned to {len(pruned_paradigms)} paradigms after removing subsumed stems.")
return sort_paradigms(pruned_paradigms)
### -----------------------------
### 6. Sort paradigms by size
### -----------------------------
def sort_paradigms(paradigms):
"""
Primary: log(len(stems)) * log(len(suffixes)) (DESC)
Ties: (-len(stems), -len(suffixes), lexicographic stems, lexicographic suffix heads)
"""
def score(p):
stems, suffixes = p
if stems and suffixes:
return math.log(len(stems)) * math.log(len(suffixes))
return 0.0
def tie_key(p):
stems, suffixes = p
sfx_heads = []
for s in suffixes:
sfx_heads.append(s[0] if isinstance(s, tuple) else s)
return (-len(stems), -len(suffixes),
" ".join(sorted(stems)),
" ".join(sorted(sfx_heads)))
return sorted(paradigms, key=lambda p: (-score(p), tie_key(p)))
def sort_paradigms_by_suffix_count(paradigms):
def score(p):
stem_count = len(p[0])
suffix_count = len(p[1])
if stem_count > 0 and suffix_count > 0:
return suffix_count
return 0
return sorted(paradigms, key=score, reverse=True)
def nest_suffixes_from_paradigms(paradigms):
print("[7/7] Nesting suffixes based on reusable paradigms...")
suffix_set_index = {frozenset(suffixes): True for _, suffixes in paradigms}
nested_paradigms = []
for stems, suffixes in paradigms:
suffixes_list = list(suffixes)
nested_suffixes = set()
used = set()
# deterministic nested pairing
for i, s1 in enumerate(sorted(suffixes_list)):
for j, s2 in enumerate(sorted(suffixes_list)):
if i == j or s2 in used or not isinstance(s1, str) or not isinstance(s2, str):
continue
if s2.startswith(s1) and s1 != '':
remainder = s2[len(s1):]
if remainder and frozenset({'', remainder}) in suffix_set_index:
nested_suffixes.add((s1, frozenset({'', remainder})))
used.add(s2)
used.add(s1)
break
for s in suffixes_list:
if s not in used:
nested_suffixes.add(s)
nested_paradigms.append((set(stems), nested_suffixes))
print(f"✅ Nested structure created for {len(nested_paradigms)} paradigms.")
return sort_paradigms(nested_paradigms)
def refine_nested_stem_conflicts(paradigms):
"""
Remove stems from higher-ranked paradigms if they are fully explained by nested structures
in lower-ranked paradigms.
Args:
paradigms: list of (stem_set, suffix_set), where suffix_set may contain nested (str, frozenset) tuples
Returns:
Refined list of paradigms with redundant derived stems removed
"""
refined_paradigms = paradigms[:]
all_suffix_sets = {frozenset(suffixes) for _, suffixes in paradigms}
# Build a mapping from nested suffix sets to their parent prefixes
derived_stems = set()
for stems, suffixes in paradigms:
for sfx in suffixes:
if isinstance(sfx, tuple):
base, nested_suffixes = sfx
if frozenset(nested_suffixes) in all_suffix_sets:
for stem in stems:
derived_stems.add(stem + base)
# Remove derived stems from paradigms with simple suffix sets (like ['', 's'])
updated_paradigms = []
for stems, suffixes in refined_paradigms:
cleaned_stems = stems - derived_stems
updated_paradigms.append((cleaned_stems, suffixes))
print(f"✅ Removed {len(derived_stems)} derived stems explained by nested paradigms.")
return updated_paradigms
### -----------------------------
### 7. Segment word based on ranked paradigms
### -----------------------------
def recursive_fallback(word, suffix_set):
for suffix in sorted(suffix_set, key=lambda s: -len(s)):
if suffix and word.endswith(suffix):
stem_candidate = word[:-len(suffix)]
rest = recursive_fallback(stem_candidate, suffix_set)
return rest + [suffix]
return [word] # fallback to whole word if nothing matches
### -----------------------------
### Main runner
### -----------------------------
def run_paradigm_extraction(vocab, min_shared_stems=2, min_suffixes=2, enrich_suffix_sets=True):
start = time.time()
stem_to_suffixes = extract_stem_suffix_pairs(vocab)
paradigms = group_stems_by_suffixes(stem_to_suffixes, min_shared_stems, min_suffixes)
paradigms = stem_set_expansion(paradigms, stem_to_suffixes)
paradigms = sort_paradigms(paradigms)
paradigms = prune_subsumed_stems(paradigms)
paradigms = sort_paradigms(paradigms)
paradigms = nest_suffixes_from_paradigms(paradigms)
paradigms = refine_nested_stem_conflicts(paradigms)
paradigms = sort_paradigms(paradigms)
if enrich_suffix_sets:
print("[4/7] Expanding suffix sets based on partial compatibility...")
paradigms = suffix_set_expansion(paradigms)
paradigms = sort_paradigms(paradigms)
paradigms = prune_subsumed_stems(paradigms)
paradigms = sort_paradigms(paradigms)
'''# Fallback paradigm for unassigned full words
vocab_words = set(vocab)
assigned_words = set()
for stems, suffixes in paradigms:
for stem in stems:
for suffix in suffixes:
if isinstance(suffix, tuple):
base, _ = suffix
assigned_words.add(stem + base)
else:
assigned_words.add(stem + suffix)
unassigned_words = vocab_words - assigned_words
if unassigned_words:
print(f"✅ {len(unassigned_words)} full words were not assigned to any paradigm, added fallback paradigm.")
paradigms.append((set(unassigned_words), frozenset({""})))
paradigms = sort_paradigms(paradigms)'''
print(f"\n✅ Extracted {len(paradigms)} paradigms.")
print(f"⏱️ Finished in {time.time() - start:.2f} seconds.")
return paradigms
def segment_word_from_nested_paradigms(word, paradigms, fallback=True, top_k=300):
"""
Segment a word based on nested paradigms with optional fallback.
Parameters:
word (str): The word to segment.
paradigms (list): A list of tuples (stems, suffixes) with optional nesting.
fallback (bool): Whether to fall back on longest suffix match from top_k paradigms.
top_k (int): Number of top paradigms to consider in fallback.
Returns:
List[str]: Segmented pieces of the word.
"""
def match_suffixes(suffixes, remainder):
"""Recursive helper to match nested suffix structures."""
for suffix in suffixes:
if isinstance(suffix, tuple):
base, nested = suffix
if remainder.startswith(base):
sub = remainder[len(base):]
nested_result = match_suffixes(nested, sub)
if nested_result is not None:
return [base] + nested_result
elif remainder == suffix:
return [suffix] if suffix else []
return None
# First pass: try full nested match
for stems, suffixes in paradigms:
for stem in stems:
if word.startswith(stem):
remainder = word[len(stem):]
matched_suffix = match_suffixes(suffixes, remainder)
if matched_suffix is not None:
return [stem] + matched_suffix
# Fallback strategy: longest suffix among top_k paradigms
if fallback:
seen_suffixes = set()
def collect_suffixes(suffixes):
for s in suffixes:
if isinstance(s, tuple):
seen_suffixes.add(s[0])
collect_suffixes(s[1])
else:
seen_suffixes.add(s)
for _, suffixes in paradigms[:top_k]:
collect_suffixes(suffixes)
# Try matching the longest suffix first
for suffix in sorted(seen_suffixes, key=lambda s: -len(s)):
if suffix and word.endswith(suffix):
stem = word[:-len(suffix)]
return [stem, suffix]
return [word]
return [word]
def segment_word_from_paradigms(word, paradigms, top_k=20):
"""
Simpler fallback-only version: match longest suffix among top_k paradigms.
Parameters:
word (str): Word to segment.
paradigms (list): Paradigm structures.
top_k (int): How many paradigms to consider.
Returns:
List[str]: Segmentation result.
"""
candidates = paradigms[:top_k]
best_split = None
for stems, suffixes in candidates:
for suffix in sorted(suffixes, key=lambda s: -len(s) if isinstance(s, str) else -len(s[0])):
if isinstance(suffix, tuple):
suffix = suffix[0] # ignore nested for fallback
if word.endswith(suffix):
stem_candidate = word[:-len(suffix)] if suffix else word
if stem_candidate in stems:
split = [stem_candidate, suffix] if suffix else [stem_candidate]
if best_split is None or len(suffix) > len(best_split[-1]):
best_split = split
return best_split or [word]

50851
paradigms.json Normal file
View File

File diff suppressed because it is too large Load Diff

6
preprocess_config.json Normal file
View File

@@ -0,0 +1,6 @@
{
"lowercase": true,
"separate_apostrophes": false,
"separate_digits": true,
"separate_punctuation": true
}

26
preprocessing.py Normal file
View File

@@ -0,0 +1,26 @@
# preprocessing.py
import re
class Preprocessor:
def __init__(self, lowercase=False, separate_apostrophes=True, separate_digits=True, separate_punctuation=True):
self.lowercase = lowercase
self.separate_apostrophes = separate_apostrophes
self.separate_punctuation = separate_punctuation
self.separate_digits = separate_digits
def preprocess(self, line: str) -> str:
if self.lowercase:
line = line.lower()
if self.separate_apostrophes:
# Add spaces around apostrophes
line = re.sub(r"([’'`])", r" \1 ", line)
# Add spaces around punctuation (except alphanumeric and apostrophes)
if self.separate_punctuation:
line = re.sub(r"([^A-Za-z0-9\s’'`])", r" \1 ", line)
if self.separate_digits:
line = re.sub(r"(\d)", r" \1 ", line)
# Normalize whitespace
line = re.sub(r"\s+", " ", line)
return line.strip()

13
revision_info.json Normal file
View File

@@ -0,0 +1,13 @@
{
"kind": "main",
"name": "main",
"checkpoint_name": "epoch_9",
"family": "epoch_end",
"epoch_index": 9,
"notes": [
"Clean staging bundle for the HF main branch.",
"This matches the packaged top-level model checkpoint."
],
"source_checkpoint_dir": "/home/achille.fusco/pr_baby_lm/BabyLM_2026_ENH/04-experiments/model_gpt2_ParFindFast_eng_zho_BD_budget16k_zhchildes_v1.0/checkpoints/epoch_9",
"target_bundle_dir": "/home/achille.fusco/pr_baby_lm/BabyLM_2026_ENH/03-models/gpt2_parfind_en_zh_equal_tokfix/main_bundle"
}

30
special_tokens_map.json Normal file
View File

@@ -0,0 +1,30 @@
{
"bos_token": {
"content": "<s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"eos_token": {
"content": "</s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<pad>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"unk_token": {
"content": "<unk>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
}
}

32274
tokenizer.json Normal file
View File

File diff suppressed because it is too large Load Diff

602
tokenizer.py Normal file
View File

@@ -0,0 +1,602 @@
# tokenizer.py
# Enhanced Paradigm-based segmenter with configurable features:
# - Word boundary tokens
# - Null suffixes as tokens
# - Paradigm-specific suffixes and roots
from collections import OrderedDict
from pathlib import Path
from typing import List, Tuple, Optional
import os, json, re
from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast
try:
from .boundary_discovery import anchor_sequences
except ImportError:
from boundary_discovery import anchor_sequences
def _deserialize_suffixes_from_json(sfx_list):
out = set()
for item in sfx_list:
if isinstance(item, list):
# JSON nested: [base, nested_list]
base, nested = item
out.add((base, frozenset(nested)))
else:
out.add(item) # plain string like "", "ing", "s"
return out
def _load_paradigms_any(path):
import json
with open(path, "r", encoding="utf-8") as f:
payload = json.load(f)
# Case A: new schema with top-level dict {"paradigms": [...]}
if isinstance(payload, dict) and "paradigms" in payload:
paradigms = []
for p in payload["paradigms"]:
stems = set(p["stems"])
suffixes = _deserialize_suffixes_from_json(p["suffixes"])
paradigms.append((stems, suffixes))
meta = payload.get("meta", {})
return paradigms, meta
# Case B: older "list of pairs" JSON [[stems, suffixes], ...]
if isinstance(payload, list) and payload and isinstance(payload[0], list):
paradigms = []
for stems, suffixes in payload:
stems = set(stems)
# suffixes may be ["", ["er", ["", "s"]], "ing"] or already strings
norm = _deserialize_suffixes_from_json(suffixes)
paradigms.append((stems, norm))
return paradigms, {}
# Case C: already python-native structure (rare if not using JSON)
if isinstance(payload, list) and payload and isinstance(payload[0], (list, tuple)) and len(payload[0]) == 2:
return payload, {}
raise ValueError("Unrecognized paradigms.json format")
# ----------------------------
# Enhanced Paradigm-based segmenter
# ----------------------------
class EnhancedParadigmFinderSegmenter:
def __init__(self, paradigms, config):
self.paradigms = paradigms
self.config = config
self.lowercase = config.get("lowercase", True)
self.space_punct = config.get("space_punct", True)
self.use_word_boundaries = config.get("use_word_boundaries", False)
self.null_suffixes_as_tokens = config.get("null_suffixes_as_tokens", False)
self.paradigm_specific_suffixes = config.get("paradigm_specific_suffixes", False)
self.paradigm_specific_roots = config.get("paradigm_specific_roots", False)
self.word_boundary_token = config.get("word_boundary_token", "▁")
self.null_suffix_token = config.get("null_suffix_token", "ε")
self.paradigm_token_format = config.get("paradigm_token_format", "{token}_p{paradigm_idx}")
self.fallback_mode = config.get("fallback_mode", "none")
self.boundaries_discovery = config.get("boundaries_discovery", False)
self.boundary_discovery_mode = config.get("boundary_discovery_mode", "space_free_only")
self.boundary_space_marker = config.get("boundary_space_marker", "_")
self.boundary_min_sequence_length = config.get("boundary_min_sequence_length", 2)
self.segment_cache_size = max(0, int(config.get("segment_cache_size", 200000)))
self._segment_cache = OrderedDict() if self.segment_cache_size > 0 else None
self.space_free_lexicon_meta = config.get("space_free_lexicon", {})
self.language_zero_morphemes = set(config.get("language_zero_morphemes", {}).values())
self._space_free_candidates_by_initial = self._build_space_free_candidates(
self.space_free_lexicon_meta
)
self._candidates_by_initial = {}
for p_idx, (stems, suffixes) in enumerate(self.paradigms):
for stem in stems:
if not stem:
continue
initial = stem[0]
if initial not in self._candidates_by_initial:
self._candidates_by_initial[initial] = {}
self._candidates_by_initial[initial].setdefault(p_idx, []).append((stem, suffixes))
for initial, paradigms_by_rank in self._candidates_by_initial.items():
ordered = []
for p_idx in sorted(paradigms_by_rank):
stems_for_rank = sorted(
paradigms_by_rank[p_idx],
key=lambda item: (-len(item[0]), item[0]),
)
ordered.append((p_idx, stems_for_rank))
self._candidates_by_initial[initial] = ordered
if self.fallback_mode not in {"none", "suffix"}:
raise ValueError(f"Unsupported fallback_mode: {self.fallback_mode}")
if self.boundary_discovery_mode not in {"space_free_only", "all"}:
raise ValueError(f"Unsupported boundary_discovery_mode: {self.boundary_discovery_mode}")
@staticmethod
def _is_han_char(ch: str) -> bool:
return (
"\u3400" <= ch <= "\u4dbf"
or "\u4e00" <= ch <= "\u9fff"
or "\uf900" <= ch <= "\ufaff"
)
def _contains_han(self, text: str) -> bool:
return any(self._is_han_char(ch) for ch in text)
def _is_zero_suffix_marker(self, suffix) -> bool:
return isinstance(suffix, str) and suffix in self.language_zero_morphemes
def _build_space_free_candidates(self, lexicon_meta):
candidates_by_initial = {}
if not isinstance(lexicon_meta, dict):
return candidates_by_initial
languages = lexicon_meta.get("languages", {})
for lang_meta in languages.values():
for token in lang_meta.get("tokens", []):
if not token or any(ch.isspace() for ch in token):
continue
initial = token[0]
candidates_by_initial.setdefault(initial, set()).add(token)
for initial, tokens in list(candidates_by_initial.items()):
candidates_by_initial[initial] = sorted(tokens, key=lambda tok: (-len(tok), tok))
return candidates_by_initial
def _segment_cache_get(self, word: str, fallback: bool, top_k: int) -> Optional[List[str]]:
if self._segment_cache is None:
return None
key = (word, fallback, top_k)
cached = self._segment_cache.get(key)
if cached is None:
return None
self._segment_cache.move_to_end(key)
return list(cached)
def _segment_cache_put(self, word: str, fallback: bool, top_k: int, pieces: List[str]) -> List[str]:
if self._segment_cache is None:
return pieces
key = (word, fallback, top_k)
self._segment_cache[key] = tuple(pieces)
self._segment_cache.move_to_end(key)
if len(self._segment_cache) > self.segment_cache_size:
self._segment_cache.popitem(last=False)
return pieces
def _format_token(self, token: str, paradigm_idx: Optional[int], apply_label: bool) -> str:
if not apply_label or paradigm_idx is None or not token:
return token
return self.paradigm_token_format.format(token=token, paradigm_idx=paradigm_idx)
def _match_suffixes(self, suffixes, remainder: str, paradigm_idx: int) -> List[List[str]]:
matches = []
for suffix in sorted(
suffixes,
key=lambda s: (
0 if isinstance(s, str) else 1,
-(0 if self._is_zero_suffix_marker(s) else (len(s) if isinstance(s, str) else len(s[0]))),
"" if self._is_zero_suffix_marker(s) else (s if isinstance(s, str) else s[0]),
),
):
if isinstance(suffix, (tuple, list)):
base, nested = suffix
if remainder.startswith(base):
sub = remainder[len(base):]
for nested_match in self._match_suffixes(nested, sub, paradigm_idx):
piece = self._format_token(
base,
paradigm_idx,
apply_label=self.paradigm_specific_suffixes,
)
if piece:
matches.append([piece] + nested_match)
else:
matches.append(nested_match)
elif self._is_zero_suffix_marker(suffix) and remainder == "":
if self.null_suffixes_as_tokens:
matches.append([
self._format_token(
self.null_suffix_token,
paradigm_idx,
apply_label=self.paradigm_specific_suffixes,
)
])
else:
matches.append([])
elif remainder == suffix:
if suffix:
matches.append([
self._format_token(
suffix,
paradigm_idx,
apply_label=self.paradigm_specific_suffixes,
)
])
elif self.null_suffixes_as_tokens:
matches.append([
self._format_token(
self.null_suffix_token,
paradigm_idx,
apply_label=self.paradigm_specific_suffixes,
)
])
else:
matches.append([])
matches.sort(key=lambda parts: (-len(parts), tuple(parts)))
return matches
def _best_full_match(self, word: str) -> Optional[List[str]]:
initial_candidates = self._candidates_by_initial.get(word[0], []) if word else []
for p_idx, stems_for_rank in initial_candidates:
best_match = None
best_score = None
for stem, suffixes in stems_for_rank:
if not word.startswith(stem):
continue
remainder = word[len(stem):]
suffix_matches = self._match_suffixes(suffixes, remainder, p_idx)
if not suffix_matches:
continue
root = self._format_token(
stem,
p_idx,
apply_label=self.paradigm_specific_roots,
)
for suffix_parts in suffix_matches:
score = (len(stem), len(suffix_parts))
if best_score is None or score > best_score:
best_score = score
best_match = [root] + suffix_parts
if best_match is not None:
return best_match
return None
def _preprocess(self, text: str) -> str:
s = text
if self.lowercase:
s = s.lower()
if self.space_punct:
s = re.sub(r"([^\w\s'])", r" \1 ", s)
s = re.sub(r"\s+", " ", s).strip()
return s
def _prepare_units(self, raw_text: str) -> List[str]:
if self._space_free_candidates_by_initial and self._contains_han(raw_text):
return self._preprocess(raw_text).split()
if not self.boundaries_discovery:
return self._preprocess(raw_text).split()
if self.boundary_discovery_mode == "space_free_only" and any(ch.isspace() for ch in raw_text.strip()):
return self._preprocess(raw_text).split()
return anchor_sequences(
raw_text.lower() if self.lowercase else raw_text,
space_marker=self.boundary_space_marker,
min_sequence_length=self.boundary_min_sequence_length,
)
def _segment_space_free_word(self, word: str) -> List[str]:
pieces = []
idx = 0
while idx < len(word):
initial = word[idx]
best = None
for candidate in self._space_free_candidates_by_initial.get(initial, []):
if word.startswith(candidate, idx):
best = candidate
break
if best is None:
best = initial
pieces.append(best)
idx += len(best)
return pieces
def _segment_word(self, word: str, fallback=True, top_k=20) -> List[str]:
"""Enhanced segmentation with deterministic paradigm selection."""
cached = self._segment_cache_get(word, fallback, top_k)
if cached is not None:
return cached
if self._space_free_candidates_by_initial and self._contains_han(word):
return self._segment_cache_put(word, fallback, top_k, self._segment_space_free_word(word))
full_match = self._best_full_match(word)
if full_match is not None:
return self._segment_cache_put(word, fallback, top_k, full_match)
if fallback and self.fallback_mode == "suffix":
candidates = self.paradigms[:top_k]
longest = ""
def collect_flat(sfx):
for s in sfx:
if isinstance(s, (tuple, list)):
yield s[0]
yield from collect_flat(s[1])
else:
yield s
for _, suffixes in candidates:
for suffix in collect_flat(suffixes):
if word.endswith(suffix) and len(suffix) > len(longest):
longest = suffix
if longest:
stem = word[:-len(longest)]
return self._segment_cache_put(word, fallback, top_k, [stem, longest])
return self._segment_cache_put(word, fallback, top_k, [word])
def segment_to_tokens(self, raw_text: str, fallback=True, top_k=20) -> List[str]:
words = self._prepare_units(raw_text)
segmented = []
for word_idx, word in enumerate(words):
if self.use_word_boundaries and word_idx > 0:
segmented.append(self.word_boundary_token)
segmented.extend(self._segment_word(word, fallback=fallback, top_k=top_k))
return segmented
def segment_with_alignment(self, raw_text: str) -> Tuple[str, List[Optional[int]]]:
"""
Return the labeled segmentation string used by both training and inference.
The alignment list is kept as a placeholder because the wrapper currently
delegates offsets to the fast tokenizer directly.
"""
segmented_tokens = self.segment_to_tokens(raw_text, fallback=True)
segmented_text = " ".join(segmented_tokens)
return segmented_text, [None] * len(segmented_text)
# ----------------------------
# Offset remapping helper
# ----------------------------
def remap_offsets_to_raw(offsets: List[Tuple[int,int]], pre2raw: List[Optional[int]]) -> List[Tuple[int,int]]:
mapped = []
L = len(pre2raw)
for s,e in offsets:
s = max(0, min(s, L)); e = max(0, min(e, L))
rs = re_ = None
t = s
while t < e and rs is None:
if pre2raw[t] is not None: rs = pre2raw[t]
t += 1
t = e - 1
while t >= s and re_ is None:
if pre2raw[t] is not None: re_ = pre2raw[t] + 1
t -= 1
mapped.append((rs if rs is not None else 0, re_ if re_ is not None else 0))
return mapped
# ----------------------------
# Public wrapper
# ----------------------------
class EnhancedParadigmTokenizerWrapper(PreTrainedTokenizerFast):
slow_tokenizer_class = None
@staticmethod
def _resolve_assets_dir(name_or_path, tok_file) -> Optional[str]:
candidates = []
if tok_file:
candidates.append(Path(tok_file).parent)
if name_or_path and os.path.isdir(name_or_path):
candidates.append(Path(name_or_path).resolve())
module_dir = Path(__file__).resolve().parent
commit_hash = module_dir.name
repo_id = name_or_path if isinstance(name_or_path, str) else None
if repo_id and "/" in repo_id:
repo_cache_name = f"models--{repo_id.replace('/', '--')}"
for parent in module_dir.parents:
candidates.append(parent / "hub" / repo_cache_name / "snapshots" / commit_hash)
for parent in module_dir.parents:
hub_root = parent / "hub"
if not hub_root.is_dir():
continue
candidates.extend(hub_root.glob(f"models--*--*/snapshots/{commit_hash}"))
seen = set()
for candidate in candidates:
candidate = Path(candidate)
key = str(candidate)
if key in seen:
continue
seen.add(key)
if (candidate / "tokenizer.json").is_file() and (candidate / "paradigms.json").is_file():
return str(candidate)
return None
@staticmethod
def _download_required_assets(name_or_path, revision, cache_dir=None, local_files_only=False) -> Optional[str]:
if not isinstance(name_or_path, str) or "/" not in name_or_path:
return None
required_files = [
"tokenizer.json",
"paradigms.json",
"preprocess_config.json",
"tokenizer_config.json",
]
downloaded = []
for filename in required_files:
try:
downloaded.append(
hf_hub_download(
repo_id=name_or_path,
filename=filename,
revision=revision,
cache_dir=cache_dir,
local_files_only=local_files_only,
)
)
except Exception:
if filename in {"paradigms.json", "tokenizer.json"}:
return None
snapshot_candidates = []
for path in downloaded:
candidate = Path(path).parent
snapshot_candidates.append(candidate)
for parent in Path(path).parents:
snapshot_candidates.append(parent / "snapshots" / str(revision))
seen = set()
for candidate in snapshot_candidates:
candidate = Path(candidate)
key = str(candidate)
if key in seen:
continue
seen.add(key)
if (candidate / "tokenizer.json").is_file() and (candidate / "paradigms.json").is_file():
return str(candidate)
return None
def __init__(self, *args, **kwargs):
# Ensure fast tokenizer is loaded directly (no slow->fast conversion)
name_or_path = kwargs.get("name_or_path", None)
if name_or_path is None and len(args) > 0 and isinstance(args[0], str):
name_or_path = args[0]
cache_dir = kwargs.get("cache_dir")
local_files_only = bool(kwargs.get("local_files_only", False))
if "tokenizer_file" not in kwargs and "tokenizer_object" not in kwargs and name_or_path is not None:
tf = os.path.join(name_or_path, "tokenizer.json")
if os.path.isfile(tf):
kwargs["tokenizer_file"] = tf
else:
commit_hash = Path(__file__).resolve().parent.name
downloaded_assets_dir = self._download_required_assets(
name_or_path=name_or_path,
revision=commit_hash,
cache_dir=cache_dir,
local_files_only=local_files_only,
)
if downloaded_assets_dir is None and local_files_only:
downloaded_assets_dir = self._download_required_assets(
name_or_path=name_or_path,
revision=commit_hash,
cache_dir=cache_dir,
local_files_only=False,
)
if downloaded_assets_dir is not None:
kwargs["tokenizer_file"] = str(Path(downloaded_assets_dir) / "tokenizer.json")
super().__init__(*args, **kwargs)
tok_file = kwargs.get("tokenizer_file", getattr(self, "tokenizer_file", None))
hf_dir = self._resolve_assets_dir(name_or_path, tok_file)
if hf_dir is None:
commit_hash = Path(__file__).resolve().parent.name
hf_dir = self._download_required_assets(
name_or_path=name_or_path,
revision=commit_hash,
cache_dir=cache_dir,
local_files_only=local_files_only,
)
if hf_dir is None and local_files_only:
hf_dir = self._download_required_assets(
name_or_path=name_or_path,
revision=commit_hash,
cache_dir=cache_dir,
local_files_only=False,
)
if hf_dir is None:
raise FileNotFoundError(
"Could not resolve local tokenizer assets directory from tokenizer_file "
f"or name_or_path={name_or_path!r}"
)
# Load paradigms
ppath = os.path.join(hf_dir, "paradigms.json")
if not os.path.exists(ppath):
raise FileNotFoundError(f"Missing paradigms.json in {hf_dir}")
self.paradigms, self.paradigms_meta = _load_paradigms_any(ppath)
# Load configuration
self.config = {}
cpath = os.path.join(hf_dir, "tokenizer_config.json")
if os.path.exists(cpath):
with open(cpath, "r", encoding="utf-8") as f:
try:
self.config.update(json.load(f))
except json.JSONDecodeError:
pass
# Load preprocessing flags
pre_cfg = {"lowercase": True, "space_punct": True}
pre_cpath = os.path.join(hf_dir, "preprocess_config.json")
if os.path.exists(pre_cpath):
with open(pre_cpath, "r", encoding="utf-8") as f:
pre_cfg.update(json.load(f))
# Merge configs
full_config = {**pre_cfg, **self.config}
if self.paradigms_meta.get("space_free_lexicon"):
full_config["space_free_lexicon"] = self.paradigms_meta["space_free_lexicon"]
if self.paradigms_meta.get("language_zero_morphemes"):
full_config["language_zero_morphemes"] = self.paradigms_meta["language_zero_morphemes"]
self.segmenter = EnhancedParadigmFinderSegmenter(
paradigms=self.paradigms,
config=full_config,
)
def _segment_input(self, value):
if isinstance(value, str):
seg, _ = self.segmenter.segment_with_alignment(value)
return seg
if isinstance(value, (list, tuple)):
segs = []
for item in value:
if not isinstance(item, str):
raise TypeError("batched inputs must contain only strings")
seg, _ = self.segmenter.segment_with_alignment(item)
segs.append(seg)
return segs
raise TypeError("text inputs must be str or List[str]/Tuple[str]")
# ---- main entry point ----
def __call__(self, text, text_pair=None, **kwargs):
seg_text = self._segment_input(text)
if text_pair is None:
return super().__call__(seg_text, **kwargs)
seg_text_pair = self._segment_input(text_pair)
return super().__call__(seg_text, text_pair=seg_text_pair, **kwargs)
def tokenize(self, text, **kwargs):
# Intercept manual .tokenize() calls to ensure segmentation happens first
if isinstance(text, str):
return super().tokenize(self._segment_input(text), **kwargs)
elif isinstance(text, list):
# Tokenize each string separately, then flatten (matches HF behavior)
out = []
for t in text:
out.extend(super().tokenize(self._segment_input(t), **kwargs))
return out
else:
raise TypeError("tokenize() expects str or List[str]")
def encode(self, text, text_pair=None, **kwargs):
seg_text = self._segment_input(text)
if text_pair is None:
return super().encode(seg_text, **kwargs)
return super().encode(seg_text, text_pair=self._segment_input(text_pair), **kwargs)
def encode_plus(self, text, text_pair=None, **kwargs):
seg_text = self._segment_input(text)
if text_pair is None:
return super().encode_plus(seg_text, **kwargs)
return super().encode_plus(
seg_text,
text_pair=self._segment_input(text_pair),
**kwargs,
)

67
tokenizer_config.json Normal file
View File

@@ -0,0 +1,67 @@
{
"added_tokens_decoder": {
"0": {
"content": "<pad>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"1": {
"content": "<unk>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"2": {
"content": "<s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"3": {
"content": "</s>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"auto_map": {
"AutoTokenizer": [
"tokenizer.EnhancedParadigmTokenizerWrapper",
null
]
},
"bos_token": "<s>",
"boundaries_discovery": true,
"boundary_discovery_mode": "space_free_only",
"boundary_min_sequence_length": 2,
"boundary_shorter_first": true,
"boundary_space_marker": "_",
"clean_up_tokenization_spaces": false,
"eos_token": "</s>",
"extra_special_tokens": {},
"fallback_mode": "none",
"model_max_length": 1000000000000000019884624838656,
"null_suffix_token": "ε",
"null_suffixes_as_tokens": false,
"pad_token": "<pad>",
"paradigm_specific_roots": false,
"paradigm_specific_suffixes": false,
"paradigm_token_format": "{token}_p{paradigm_idx}",
"segment_cache_size": 200000,
"space_free_bootstrap_languages": [
"zho"
],
"tokenizer_class": "EnhancedParadigmTokenizerWrapper",
"unk_token": "<unk>",
"use_word_boundaries": false,
"word_boundary_token": "▁"
}