File size: 94,427 Bytes
429d18d | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 1670 1671 1672 1673 1674 1675 1676 1677 1678 1679 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709 1710 1711 1712 1713 1714 1715 1716 1717 1718 1719 1720 1721 1722 1723 1724 1725 1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812 1813 1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 1837 1838 1839 1840 1841 1842 1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865 1866 1867 1868 1869 1870 1871 1872 1873 1874 1875 1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893 1894 1895 1896 1897 1898 1899 1900 1901 1902 1903 1904 1905 1906 1907 1908 1909 1910 1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926 1927 1928 1929 1930 1931 1932 1933 1934 1935 1936 1937 1938 1939 1940 1941 1942 1943 1944 1945 1946 1947 1948 1949 1950 1951 1952 1953 1954 1955 1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 | __version__ = "1.4.8"
__author__ = "Cristiano Chesi & NeTS Lab @ IUSS (with Claude Sonnet 4.6 / Opus 4.8 fixes)"
__email__ = "cristiano.chesi@iusspavia.it"
__status__ = "Research"
__date__ = "2026-06-30"
__license__ = "MIT"
"""
MorPiece is a split-based tokenization library that incrementally chunks words into potentially meaningful morphemes. The splitting procedure consists of evaluating if the Tolerance Principle (Yang 2016) applies at every character every time an incoming word "traverses" the lexicon.
Take the word "cats": a root "trie" (c->a->t->s) and a inflectional trie (s->t->a->c) are considered.
In the default mode (token-based) "traversal" the lexicon means adding 1 to each node counter that is traversed both in the root trie and in the inflectional trie. If a path does not exists, it is initialized to 1.
The type-based "traversal" modality updates the node counter only when new types are observed.
A split between "t" and "s" is pustulated if and only if both in the root trie and in the infl trie the tolerance principle is respected, that is:
freq(t)/ln(freq(t) > freq(s) in the root trie and
freq(s)/ln(freq(s) > freq(t) in the infl trie
if this is the case, the "s" pendant (in this case just "s") is added to the root trie, under the "++" node.
At the end, all nodes that does not have a frequency above min_freq parameter are pruned.
MaxLength strategy is adopted to retrieve the tokens for each word.
Since version 1.4.* a type-based tokenization option is implemented ('type_based = true')
If the "order or acquisition" (ooa) parameter is set to True, each 100K of exposure a vocabulary is created to check splitting hypotheses postulated and the evidence needed (for research purposes)
NOTE: by default (char_coverage=0.99) the most frequent characters are added to the final vocabulary at the end of training, each in both its root ("c") and inflectional ("++c") form, up to 99% of character occurrences. This gives one token per glyph (ideal for CJK — "中文" -> "中" "++文"), works identically in the exported HuggingFace WordPiece tokenizer, and leaves the WordPiece segmentation of known words unchanged. The added tokens are baked into the saved tokenizer, so char_coverage need not be set when loading a trained tokenizer.
Examples:
import tokenizer_MorPiece as MoP
mop = MoP.MorPiece(vocab_size=vocab_size, cutoff=cutoff, min_suffix_stems=3, min_frequency=min_frequency, ooa=ooa, use_tokenizers_lib=True)
mop.train(text)
mop.save('./mop_tokenizer/tokenizer.json')
s = "test sentence"
print("Sentence to tokenize: " + s)
ids, tokens = mp.encode(s)
print(ids, tokens)
Todo:
- Multi word evaluation before splitting
Reference:
https://github.com/cristianochesi/morpiece
"""
import os
import re
import json
from collections import deque
from math import log
from tokenizers import pre_tokenizers, decoders, normalizers, Regex, processors
class MorPiece:
"""MorPiece incrementally chunks words into potentially meaningful morphemes."""
ids: list[int]
tokens: list[str]
vocab_size: int
# -------------------------------------------------------------------------
# CHILDES / CHAT speaker codes (v1.4.4)
# -------------------------------------------------------------------------
# Common line-initial CHAT speaker tiers. Each entry C becomes the atomic
# special token "*C:". This is only a SEED set for cross-run consistency:
# any speaker code actually found in the corpus — standard three-letter or
# spelled-out, e.g. "*INVESTIGATOR:" — is detected and registered on the
# fly during preprocessing, so the list need not be exhaustive.
CHILDES_SPEAKER_CODES = (
"CHI", "MOT", "FAT", "INV", "EXP", "GRA", "GRM", "GRF",
"SIS", "BRO", "SIB", "ADU", "TEA", "CAR", "BAB", "AUN",
"UNC", "COU", "FRI", "NEI", "OPE", "PAR", "UNK", "NON", "TOY",
"CHILD", "MOTHER", "FATHER", "INVESTIGATOR", "EXPERIMENTER",
"TEACHER", "ADULT", "GRANDMOTHER", "GRANDFATHER",
)
def __init__(
self,
vocab_size=30000,
min_frequency=10,
cutoff=100,
special_tokens=None,
ooa=True,
use_tokenizers_lib=True,
type_based=True,
use_speaker_tokens=False,
childes_speaker_tokens=True,
min_suffix_stems=3,
ooa_type_interval=1000,
ooa_token_interval=100000,
counter_unit='tokens',
lowercase=True,
boundaries_discovery=False,
char_coverage=0.99,
byte_fallback=True,
glue_morphemes=False,
):
"""
Parameters
----------
vocab_size : int
Maximum vocabulary size (default 30000).
min_frequency : int
Minimum token frequency to survive final pruning (default 10).
cutoff : int
Minimum mother-node frequency before TP is evaluated. Nodes below
this threshold are ignored (default 100). Acts as the principal
guard against splits at low-evidence positions.
ooa : bool
Save Order-of-Acquisition vocabulary snapshots (default True).
use_tokenizers_lib : bool
Use HuggingFace tokenizers for normalisation and pre-tokenisation.
type_based : bool
Update trie counts once per type instead of once per token.
⚠ STRUCTURAL LIMITATION for morphologically rich languages:
In type_based mode, trie counts equal the number of distinct types
sharing a path prefix — NOT token frequencies. For Italian with
~200 K types, paradigm-segment counts (e.g. "cantav-") are only
5–15 types. The TP threshold tp = m / log(m) sits at ~3–7, and
the daughter count d must exceed tp. This barely succeeds at best.
In TOKEN-BASED mode (type_based=False) the same segment sees
5 000–20 000 tokens across all its forms; the threshold sits at
~1 000–1 500 and is easily cleared. For Italian morphological
tokenization, type_based=False is strongly recommended.
use_speaker_tokens : bool
Detect and replace line-initial "A:"–"E:" labels with <speaker_X>
special tokens. Mid-word labels ("INCIDAZIONE:") are left untouched.
childes_speaker_tokens : bool
Detect line-initial CHILDES / CHAT speaker tiers ("*CHI:", "*MOT:",
"*INVESTIGATOR:" …) and keep each one as a single atomic special
token instead of letting the pre-tokeniser shred it into
"*" + "chi" + ":". Detection is pattern-based and registers any
speaker code it meets; CHILDES_SPEAKER_CODES seeds the common ones.
Default True — harmless on non-CHAT corpora, where the line-initial
"*code:" pattern simply does not occur.
min_suffix_stems : int
Minimum distinct root stems a ++ suffix must derive from (default 3).
Prunes coincidental character overlaps (e.g. "nto" from only "lento").
ooa_type_interval : int
Types between OOA snapshots in type_based mode (default 1000).
ooa_token_interval : int
Amount of `counter_unit` between OOA snapshots / incremental-cleaning
passes in TOKEN-based mode (and in --boundaries_discovery). Default
100000. Previously hard-coded to 100000 tokens.
counter_unit : str
Unit in which token-based OOA exposure is accumulated:
'tokens' (default; one increment per processed token — exactly the
previous behaviour), 'chars' (UTF-8 characters), or 'syllables'
(approximate vowel-group / per-CJK-glyph count). The syllable is the
most cross-linguistically comparable measure of input quantity
(Räsänen et al. 2019, 2021), so for EN-vs-IT parity prefer
counter_unit='syllables' over a fixed word-token interval. Applies
to the token-based and boundary-discovery snapshot schedules only;
type_based snapshots remain measured in types via ooa_type_interval.
special_tokens : list or None
Reserved tokens; defaults to the standard set.
"""
if special_tokens is None:
special_tokens = ['<unk>', '<pad>', '<s>', '</s>', '<mask>', '<sep>', '<cls>']
self.use_speaker_tokens = use_speaker_tokens
self.lowercase = lowercase
self.speaker_token_map = (
{L: f'<speaker_{L}>' for L in 'ABCDE'}
if use_speaker_tokens else {}
)
if use_speaker_tokens:
for st in self.speaker_token_map.values():
if st not in special_tokens:
special_tokens = list(special_tokens) + [st]
# --- v1.4.4: CHILDES / CHAT speaker tiers as atomic special tokens ---
self.childes_speaker_tokens = childes_speaker_tokens
if childes_speaker_tokens:
special_tokens = list(special_tokens)
for code in self.CHILDES_SPEAKER_CODES:
tok = code
if self.lowercase:
tok = code.lower()
if tok not in special_tokens:
special_tokens.append(tok)
self.special_tokens = special_tokens
self.unk_token_id = 0
self.pad_token_id = 1
self.bos_token_id = 2
self.eos_token_id = 3
self.mask_token_id = 4
self.sep_token_id = 5
self.cls_token_id = 6
self.start_of_text_symbol = '<s>'
self.reserved_keys = {'[RSX]', '##', 'IDX', '++'}
self.vocab_size = vocab_size
self.min_frequency = min_frequency
self.ooa = ooa
self.ooa_split = {}
self.ooa_type_interval = ooa_type_interval
if counter_unit not in ('tokens', 'chars', 'syllables'):
raise ValueError(
f"counter_unit must be 'tokens', 'chars' or 'syllables', "
f"got {counter_unit!r}"
)
self.ooa_token_interval = ooa_token_interval
self.counter_unit = counter_unit
self.use_tokenizers_lib = use_tokenizers_lib
self.type_based = type_based
self.min_suffix_stems = min_suffix_stems
self.suffix_stems: dict[str, set] = {}
self.roots = {'[RSX]': {}, '++': {}}
self.infls = {}
# --- v1.4.5 (--boundaries_discovery) ---------------------------------
# When True, MoP does NOT pre-tokenise on whitespace. The corpus is cut
# only at *unambiguous* edges (punctuation + line breaks = true starts
# "[[" / true ends "]]"); each anchored sequence traverses the root trie
# from its true start and the infl trie from its true end, and boundaries
# "||" are placed incrementally as soon as the bilateral TP licenses them
# (sufficiency). Spaces are demoted to an ordinary soft-cue symbol so
# word boundaries re-emerge from statistics rather than being given.
self.boundaries_discovery = boundaries_discovery
self.SPACE_MARK = '\u2581' # soft space cue (HF-friendly)
# v1.4.5: at the end of training, add the most frequent characters as
# single-character vocab tokens (root "c" + infl "++c") up to this
# fraction of character occurrences. One token per glyph — far better
# than a byte fallback for CJK. Default 0.99; set 0/None to disable.
# See _add_frequent_chars. Irrelevant at load (the tokens are saved).
self.char_coverage = char_coverage
# v1.4.7: byte-level fallback tail BENEATH char coverage. Any character
# that is not matchable in the trie (rare/unseen hanzi in particular) is
# decomposed into its UTF-8 bytes as dedicated single-byte tokens
# "<0x00>".."<0xFF>" (SentencePiece/Llama convention) instead of collapsing
# to a single shared <unk>. This guarantees two DISTINCT characters never
# produce the same token sequence -> minimal pairs stop tying -> the
# PinyinBench/HanziBench 0.0000 collapse disappears. Covered chars remain
# ONE token per glyph (char_coverage); only the uncovered tail pays the
# 3-bytes-per-CJK-glyph cost. The 256 byte tokens are added late (after
# _add_frequent_chars) so existing IDs / Latin segmentation are untouched.
self.byte_fallback = byte_fallback
self.BYTE_TOKENS = [f"<0x{b:02X}>" for b in range(256)]
self._byte_tok_set = set(self.BYTE_TOKENS)
# v1.4.8: unify the root/continuation embedding. When glue_morphemes is
# on, we do NOT mint a distinct "++X" token for every root X. Instead a
# word-internal morpheme reuses its ROOT token/embedding, and a single
# shared GLUE token "++" is emitted between concatenated morphemes:
# "superhero" -> ["super", "++", "hero"] (hero == the root "hero")
# Genuine learned suffix SURFACES are folded into the root vocab (one
# embedding per surface), and roots['++'] collapses to the single glue
# leaf. Segmentation is plain greedy-longest over the root vocab, which
# already realises "prefer the bigger token" (root "hero" beats any short
# suffix piece). "++" is meant to be learned as a composition hint.
self.glue_morphemes = glue_morphemes
self.GLUE = '++'
# v1.4.8: CJK has no concatenative inflection — its productive
# character-continuations are COMPOUNDING (root+root), not affixation.
# So for a CJK continuation we prefer root+glue and only let a "++X"
# suffix win when it is STRICTLY LONGER (a genuine multi-char bound
# morpheme). This stops a compound member like 脑 (in 电脑) from being
# mis-analysed as a bound suffix "++脑". Latin/NLD keep the inflection-
# first tie-break. Set False to treat CJK like Latin.
self.glue_cjk_prefer_root = True
self._CJK_RE = re.compile(r'[\u3400-\u9fff\uf900-\ufaff\U00020000-\U0002ffff]')
self._anchor_re = re.compile(r"[.!?;:,…·。！？；：，、（）()\[\]{}\"'«»\n\r\t]+")
self.types = {}
self.idx = 0
self.ids = []
self.tokens = []
self.cutoff = cutoff
self.num_tokens_in_corpus = 0
self._ooa_exposure = 0 # running token position for the OOA event log
self.num_chars_in_corpus = 0
self.num_chars_in_trie = 0
self.num_chars_in_optimized_trie = 0
# --- special-token splitter cache (see _special_token_splitter) ------
self._split_re = None
self._split_re_n = -1
self._special_set = set()
self.set_special_tokens(self.special_tokens)
# m_daughters is still tracked for OOA diagnostic output (it tells us
# how branchy each node is, which is useful for parameter tuning) —
# it just no longer participates in the TP decision as of 1.4.2.
self.ooa_data = [['Word', 'Piece',
'Root Trie - Mother freq', 'Root Trie - Daughter freq',
'Root Trie - Daughter # of children',
'Infl Trie - Mother freq', 'Infl Trie - Daughter freq',
'Infl Trie - Daughter # of children',
'Vocabulary size',
'Token exposure']]
self.vocab_to_id = None
self.vocab_to_freq = None
self.id_to_vocab = None
self.current_world = ""
self.node = {}
self.m_freq = 1
self.d_freq = 1
self.m_daughters = 1
if self.use_tokenizers_lib:
self._init_tokenizers_components()
# =========================================================================
# HF tokenizers initialisation
# =========================================================================
def _init_tokenizers_components(self):
self.normalizer = normalizers.Sequence([
normalizers.Lowercase() if self.lowercase else None,
normalizers.Prepend(" "),
normalizers.NFKC(),
# --- v1.4.4: two-tier punctuation strategy -----------------------
# (1) Unify typographic apostrophes (U+2019 ' , U+02BC ʼ) to ASCII
# U+0027 so Tier 2 sees a single codepoint.
normalizers.Replace(Regex("[\u2019\u02bc]"), "'"),
# (2) TIER 1 — hard-isolate every special character EXCEPT the
# apostrophe (a "special char" = not whitespace, letter, digit,
# or apostrophe). Two zero-width insertions, because
# normalizers.Replace does literal replacement (no $1 backrefs):
# 2a — space BEFORE such a char (if not already space-led)
# 2b — space AFTER such a char (if not already space-followed)
# Result: each is its own token, separated from BOTH neighbours.
# "house!!!" -> "house ! ! !"
# "..@%$£-!!!" -> ". . @ % $ £ - ! ! !"
# "(parola)" -> "( parola )"
normalizers.Replace(Regex("(?<=\\S)(?=[^\\s\\p{L}\\p{N}'])"), " "),
normalizers.Replace(Regex("(?<=[^\\s\\p{L}\\p{N}'])(?=\\S)"), " "),
# TIER 2 — the apostrophe is deliberately NOT touched here.
# It is kept word-internal (see pre-tokenizer) so the morpheme
# boundary is set by BPE statistics, not by a positional rule
# that would be linguistically wrong in some language (Italian
# elides left: l', d'; English contracts right & irregularly:
# do|n't, I|'m). Word-final apostrophes ("po'", "dogs'") are
# covered too — no letter follows, nothing to split.
normalizers.Replace(Regex("\n"), '\n '),
normalizers.Replace(Regex(" *\n"), '\n'),
])
self.pre_tokenizer = pre_tokenizers.Sequence([
# TIER 2 cont. — U+0027 is added to the letter classes (the trailing
# ' inside each [...] below) so an apostrophe flanked by letters does
# NOT break a word: "l'acqua", "don't", "qu'il" stay ONE pre-token,
# letting MoP learn the language-specific clitic/elision split.
pre_tokenizers.Split(
Regex(
"[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}']*"
"[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}']+|[^\\r\\n\\p{L}\\p{N}]?"
"[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}']+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}']*"
"| ?\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
),
behavior="isolated",
invert=False,
),
pre_tokenizers.Split(
Regex(".{1,24}"),
behavior="isolated",
invert=False,
),
])
self.decoder = decoders.Sequence([
decoders.Strip(' ', 1, 0),
decoders.Replace("\n ", "\n"),
])
self.post_processor = processors.TemplateProcessing(
single=f"{self.start_of_text_symbol} $A",
pair=f"{self.start_of_text_symbol} $A {self.start_of_text_symbol} $B",
special_tokens=[(self.start_of_text_symbol, self.bos_token_id)],
)
# =========================================================================
# Speaker-label preprocessing
# =========================================================================
_SPEAKER_RE = re.compile(r'(?m)^[ \t]*([A-Ea-e]):[ \t]*')
# v1.4.4: line-initial CHILDES / CHAT speaker tier, e.g. "*CHI:" / "*MOT:"
# / "*INVESTIGATOR:". The "*" and the code are adjacent, and the code and
# the ":" are adjacent — this is what distinguishes a real speaker tier
# from a markdown bullet such as "* word:".
# _CHILDES_SPEAKER_RE = re.compile(r'(?m)^[ \t]*\*([A-Za-z]{2,20}):')
def _register_special_token(self, tok: str) -> None:
"""Register a token as an atomic special token (RSX + special list)."""
if tok not in self.roots['[RSX]']:
self.roots['[RSX]'][tok] = {'IDX': self.idx}
self.idx += 1
if tok not in self.special_tokens:
self.special_tokens.append(tok)
# def _process_speaker_labels(self, text: str) -> str:
# """Replace line-initial speaker labels with atomic special tokens.
# Handles the generic "A:"–"E:" dialogue labels (-> <speaker_X>) and,
# when childes_speaker_tokens is on, CHILDES / CHAT tiers "*CODE:"
# (kept verbatim, uppercased). Used on BOTH the tokenizers and the
# non-tokenizers code path: it runs on the raw text, before any HF
# normalisation, and registers any speaker code it meets on the fly.
# """
# if self.childes_speaker_tokens:
# def _childes(m: re.Match) -> str:
# tok = '*' + m.group(1).upper() + ':'
# self._register_special_token(tok)
# return ' ' + tok + ' '
# text = self._CHILDES_SPEAKER_RE.sub(_childes, text)
# if not self.use_speaker_tokens:
# return text
# def _replace(m: re.Match) -> str:
# letter = m.group(1).upper()
# token = self.speaker_token_map.get(letter)
# return (token + ' ') if token else m.group(0)
# return self._SPEAKER_RE.sub(_replace, text)
# =========================================================================
# Text preprocessing
# =========================================================================
def _special_token_splitter(self):
"""Return a compiled regex that splits text on any registered special token.
Longest-token-first alternation, so a longer token wins over a shorter
prefix of it (e.g. "*INVESTIGATOR:" is preferred over a hypothetical
shorter "*INV:"). The capturing group means re.split() keeps the
special tokens as separate elements in its output.
The regex is cached and rebuilt only when the special-token set
changes — which can happen during early training, as
_process_speaker_labels registers new speaker codes on the fly.
Returns None if no special tokens are registered.
"""
specials = [t for t in self.roots['[RSX]'] if t]
if self._split_re_n != len(specials):
specials.sort(key=len, reverse=True)
self._special_set = set(specials)
self._split_re = (
re.compile('(' + '|'.join(re.escape(s) for s in specials) + ')')
if specials else None
)
self._split_re_n = len(specials)
return self._split_re
def _preprocess_text(self, text: str) -> str:
"""Full preprocessing pipeline.
Special tokens are preserved verbatim. Speaker labels are detected
and registered on the RAW text by _process_speaker_labels; the text is
then split on every registered special token, and ONLY the non-special
segments are sent through HF normalisation + pre-tokenisation. A
special token can therefore never be lowercased, NFKC-folded, or
shredded by the TIER 1 punctuation rules ("<pad>" -> "< pad >",
"*CHI:" -> "* chi :").
"""
# Speaker-label detection + on-the-fly registration runs on raw text,
# for both code paths.
# text = self._process_speaker_labels(text)
if not self.use_tokenizers_lib:
return text
split_re = self._special_token_splitter()
segments = split_re.split(text) if split_re else [text]
out = []
for seg in segments:
if not seg:
continue
if seg in self._special_set:
# Verbatim — a special token bypasses normalisation entirely;
# a special token is also a hard word boundary.
out.append(seg)
else:
normalized = self.normalizer.normalize_str(seg)
out.extend(tok for tok, _ in
self.pre_tokenizer.pre_tokenize_str(normalized))
return " ".join(out)
def _postprocess_tokens(self, token_ids: list) -> list:
if not self.use_tokenizers_lib:
return token_ids
if token_ids and token_ids[0] != self.bos_token_id:
return [self.bos_token_id] + token_ids
return token_ids
# =========================================================================
# Core trie operations
# =========================================================================
def __morsplit(self, w: str) -> None:
reversed_w = w[::-1]
for i in range(2, len(w) + 1):
self.__find_path(reversed_w[:len(w) + 2 - i], self.infls)
i_m, i_d, i_nd = self.m_freq, self.d_freq, self.m_daughters
infls_tp = self.__check_tp(self.m_freq, self.d_freq)
self.__find_path(w[:i], self.roots)
r_m, r_d, r_nd = self.m_freq, self.d_freq, self.m_daughters
roots_tp = self.__check_tp(self.m_freq, self.d_freq)
if roots_tp and infls_tp:
stem = w[:i - 1]
suffix = w[i - 1:]
if suffix not in self.suffix_stems:
self.suffix_stems[suffix] = set()
self.suffix_stems[suffix].add(stem)
if self.ooa:
key = self.current_world + " " + stem + "-" + suffix
if key in self.ooa_split:
self.ooa_split[key] += 1
else:
self.ooa_data.append([
self.current_world, stem + "-" + suffix,
r_m, r_d, r_nd, i_m, i_d, i_nd,
len(self.types), self._ooa_exposure,
])
self.ooa_split[key] = 1
stem_node = self.roots
for ch in stem:
stem_node = stem_node[ch]
if 'IDX' not in stem_node:
stem_node['IDX'] = 1
self.__build_trie(suffix, self.roots['++'])
def __build_trie(self, wordpiece: str, root: dict) -> None:
node = root
for ch in wordpiece:
if ch in node:
node[ch]['##'] += 1
else:
node[ch] = {'##': 1}
node = node[ch]
if 'IDX' not in node:
node['IDX'] = 1
@staticmethod
def __incremental_cleaning(trie, freq):
"""Remove low-frequency pendant nodes (token-based mode only)."""
queue = deque([(None, None, trie)])
while queue:
parent, key, current = queue.popleft()
if isinstance(current, dict):
if key == "[RSX]":
continue
if (key != "++") and (current.get("##", 0) <= freq):
if parent is not None:
del parent[key]
continue
for k, v in list(current.items()):
if isinstance(v, dict):
queue.append((current, k, v))
def __optimize(self, trie) -> None:
queue = deque([(None, None, trie)])
while queue:
parent, key, current = queue.popleft()
if isinstance(current, dict):
if key == "[RSX]":
continue
if (key != "++") and (current.get("##", 0) <= self.min_frequency):
if parent is not None:
del parent[key]
continue
if current.get("##", None) is not None:
del current["##"]
if "IDX" in current:
if self.idx < self.vocab_size:
current["IDX"] = self.idx
self.idx += 1
else:
del parent[key]
for k, v in list(current.items()):
if isinstance(v, dict):
queue.append((current, k, v))
self.__build_vocab_lookup()
def __find_path(self, word: str, trie: dict) -> None:
node = trie
for ch in word:
self.node = node[ch]
self.m_freq = self.d_freq
self.d_freq = node[ch]['##']
self.m_daughters = len(node[ch])
node = node[ch]
def __retrieve(self, string: str, trie: dict) -> None:
"""Greedy longest-match segmentation, WordPiece-equivalent.
glue_morphemes: matches every position against the ROOT vocab and emits
a single shared "++" glue token before each word-internal morpheme, so
the morpheme reuses its root embedding ("superhero" -> super, ++, hero).
"""
if getattr(self, 'vocab_to_id', None) is None:
self.__build_vocab_lookup()
vocab = self.vocab_to_id
unk_id = self.roots['[RSX]']['<unk>']['IDX']
glue = self.glue_morphemes
glue_id = vocab.get(self.GLUE) if glue else None
n = len(string)
pieces = [] # (token_text, token_id)
start = 0
while start < n:
match = None # (token_text, token_id, span, is_glue_root)
if glue and start > 0:
# HYBRID continuation. Find the longest genuine "++X" suffix AND
# the longest plain root at this position, then prefer the suffix
# UNLESS a root exhausts a strictly longer chunk (a compound part,
# e.g. "super|hero"). Keeps inflection "++ing" distinct from the
# root "ing"; reuses the root embedding for compound members.
best_suf = best_root = None # each: (token, tid, span)
end = n
while end > start:
sub = string[start:end]
span = end - start
if best_suf is None:
t = vocab.get('++' + sub)
if t is not None:
best_suf = ('++' + sub, t, span)
if best_root is None:
t = vocab.get(sub)
if t is not None:
best_root = (sub, t, span)
if best_suf is not None and best_root is not None:
break
end -= 1
if best_root is not None:
if best_suf is None:
prefer_root = True
elif (self.glue_cjk_prefer_root
and self._CJK_RE.match(best_root[0])):
prefer_root = best_root[2] >= best_suf[2] # CJK: ties -> root
else:
prefer_root = best_root[2] > best_suf[2] # Latin: ties -> suffix
if prefer_root:
match = (best_root[0], best_root[1], best_root[2], True)
else:
match = (best_suf[0], best_suf[1], best_suf[2], False)
elif best_suf is not None:
match = (best_suf[0], best_suf[1], best_suf[2], False)
else:
# word-initial (or legacy mode): plain greedy-longest match.
end = n
while end > start:
sub = string[start:end]
cand = sub if (glue or start == 0) else ('++' + sub)
tid = vocab.get(cand)
if tid is not None:
match = (cand, tid, end - start, False)
break
end -= 1
if match is None:
# No trie match for even a single character at this position.
# v1.4.7: emit that ONE character as its UTF-8 byte tokens and
# keep going, instead of collapsing the whole word to <unk>.
# This is what stops distinct rare hanzi from producing identical
# (all-<unk>) token sequences in minimal pairs.
if self.byte_fallback:
for b in string[start].encode('utf-8'):
bt = self.BYTE_TOKENS[b]
bid = vocab.get(bt)
pieces.append((bt, bid) if bid is not None
else ('<unk>', unk_id))
start += 1 # advance past exactly one character
continue
# legacy behaviour (byte_fallback=False): whole word -> <unk>
self.tokens[-1] = '<unk>'
self.ids[-1] = unk_id
return
if match[3] and glue_id is not None: # compound part -> emit glue first
pieces.append((self.GLUE, glue_id))
pieces.append((match[0], match[1]))
start += match[2]
# prepare_encoding() seeded a placeholder slot (tokens[-1]="", ids[-1]=0)
if not pieces:
self.tokens[-1] = '<unk>'
self.ids[-1] = unk_id
return
self.tokens[-1] = pieces[0][0]
self.ids[-1] = pieces[0][1]
for tok, tid in pieces[1:]:
self.tokens.append(tok)
self.ids.append(tid)
def __check_tp(self, m, d) -> bool:
"""
Tolerance Principle check (Yang 2016) — v1.4.2 form.
Returns True iff:
(a) mother frequency m exceeds the cutoff, AND
(b) daughter count d strictly exceeds the tolerance threshold
tp = m / log(m), AND
(c) d != m (the daughter does not exhaust the mother — i.e.,
there is at least some alternative continuation, otherwise
there is nothing to be productive *over*).
The branching-factor filter (`nd < bf`) of versions ≤1.4.1 has
been removed: high-branching nodes are exactly the inflectional
decision points where productivity should be evaluated, not
suppressed. The cutoff plus the bilateral root×infl requirement
plus min_suffix_stems already guards against spurious splits at
low-evidence prefixes.
"""
if m <= self.cutoff:
return False
tp = m / log(m)
return (m != d) and (d > tp)
def _add_frequent_chars(self) -> None:
"""Ensure the most frequent characters are in the vocabulary as single
tokens, up to `char_coverage` of character occurrences (default 0.99).
Rationale (v1.4.5): rather than a byte fallback (which shreds each CJK
glyph into 3 UTF-8 byte tokens), we add the actual *characters* — one
token per glyph. For every frequent character we make sure BOTH forms
exist:
• root form "c" (word-initial)
• inflectional form "++c" (word-internal / continuation)
so e.g. "中文你好" tokenises as ["中", "++文", "++你", "++好"] — and,
because these are ordinary single-character vocab entries, the SAME
thing happens in the exported HuggingFace WordPiece tokenizer (no model
change, no fragmentation).
Character frequencies are token-weighted over the (alphabetic) word
types, so coverage is measured over real character occurrences. This
runs AFTER __optimize, so it only ADDS vocabulary entries and never
changes how already-known words are segmented (a word that contains
none of the newly-added characters is untouched; for a word that does,
the deeper greedy match still wins, so Latin segmentation — and the
SIGMORPHON behaviour — is unchanged). The added tokens are written into
self.roots, so they are saved with the tokenizer and present on reload.
Controlled by `char_coverage` (default 0.99): set to 0/None to disable.
Characters beyond the coverage threshold (the rare <1% tail) are left
out and remain "<unk>".
"""
cov = self.char_coverage
if not cov or cov <= 0:
return
if not self.types:
# no per-type counts available (e.g. boundaries_discovery) -> nothing to do
return
# token-weighted character frequencies over the alphabetic word types
char_freq: dict[str, int] = {}
for word, count in self.types.items():
for ch in word:
char_freq[ch] = char_freq.get(ch, 0) + count
total = sum(char_freq.values())
if total == 0:
return
target = cov * total
ranked = sorted(char_freq.items(), key=lambda kv: kv[1], reverse=True)
def _ensure_idx(container: dict, key: str) -> bool:
"""Give `container[key]` an IDX if it lacks one; return True if added."""
node = container.get(key)
if isinstance(node, dict):
if 'IDX' in node:
return False
node['IDX'] = self.idx
else:
container[key] = {'IDX': self.idx}
self.idx += 1
return True
added_root = added_infl = n_chars = 0
cum = 0
for ch, f in ranked:
if cum >= target:
break
cum += f
n_chars += 1
if ch not in self.vocab_to_id: # root form "c"
added_root += _ensure_idx(self.roots, ch)
if ('++' + ch) not in self.vocab_to_id: # infl form "++c"
added_infl += _ensure_idx(self.roots['++'], ch)
self.__build_vocab_lookup()
if added_root or added_infl:
print(f" Added frequent-char coverage to {cov:.0%}: "
f"+{added_root} root, +{added_infl} ++ tokens "
f"(covering {n_chars} distinct chars).")
def _add_byte_fallback_tokens(self) -> None:
"""Register the 256 single-byte fallback tokens "<0x00>".."<0xFF>".
Added AFTER char coverage so they occupy the highest ids and never shift
an existing token's id or the segmentation of any in-vocab word. They are
written into self.roots (root form only — byte tokens carry no "++"
continuation prefix, matching the SentencePiece/HF ByteFallback
convention), so they persist through save_pretrained / from_pretrained and
are present in vocab_to_id for __retrieve to emit.
"""
if not self.byte_fallback:
return
if getattr(self, 'vocab_to_id', None) is None:
self.__build_vocab_lookup()
added = 0
for bt in self.BYTE_TOKENS:
if bt not in self.vocab_to_id and bt not in self.roots:
self.roots[bt] = {'IDX': self.idx}
self.idx += 1
added += 1
if added:
self.__build_vocab_lookup()
print(f" Added byte-level fallback: +{added} byte tokens "
f"(<0x00>..<0xFF>) — no character will collapse to <unk>.")
def _add_root_continuations(self) -> None:
"""Collapse the root/suffix distinction: register every ROOT token "X"
also as a word-internal continuation "++X" (with its own distinct id),
so compounds segment into their parts ("snow"+"++ball") — exactly like a
learned suffix, and identically in the native tokenizer and the exported
HuggingFace WordPiece model.
A distinct id (not a re-use of X's id) is REQUIRED: PreTrainedTokenizerFast
rebuilds its vocabulary by inverting id->token on load and silently drops
tokens that share an id, which corrupts tokenization. The "++X"
continuation is therefore a first-class token that lives in the inflection
trie, so it persists through save_pretrained / from_pretrained and needs no
special handling in __retrieve or save_HF.
Learned inflectional suffixes are never touched — an existing "++X" is kept
as-is — so inflectional segmentation is unchanged.
"""
if not getattr(self, 'vocab_to_id', None):
self.__build_vocab_lookup()
specials = {'<unk>', '<pad>', '<s>', '</s>', '<mask>'}
infl_trie = self.roots['++']
added = 0
for token in list(self.vocab_to_id.keys()):
if (not token) or token.startswith('++') or token in specials:
continue
if ('++' + token) in self.vocab_to_id: # keep learned suffixes
continue
node = infl_trie
for ch in token:
nxt = node.get(ch)
if not isinstance(nxt, dict):
nxt = {}
node[ch] = nxt
node = nxt
if 'IDX' not in node:
node['IDX'] = self.idx
self.idx += 1
added += 1
if added:
self.__build_vocab_lookup()
print(f" Collapsed root/suffix distinction: +{added} ++root "
f"continuation tokens (compounds now segment; "
f"vocab = {self.get_vocab_size()}).")
def _finalize_glue_vocab(self) -> None:
"""glue_morphemes mode. KEEP every Tolerance-discovered inflectional
suffix as its own "++X" token (an inflection "++ing" is a different
morpheme from a root "ing" as in "ingest" — they must not share an
embedding). We only:
(1) do NOT run _add_root_continuations (so no redundant "++root"
duplicate is minted for every root), and
(2) add ONE shared glue token "++".
At encode time __retrieve prefers a genuine "++X" suffix, and re-checks
the root trie only to override with "++" + root when a root match is
STRICTLY longer (a compound part: "superhero" -> super ++ hero, hero ==
the root). Purely additive — no id is moved, so no re-index is needed.
"""
if '++' not in self.roots or not isinstance(self.roots['++'], dict):
self.roots['++'] = {}
if 'IDX' not in self.roots['++']:
self.roots['++']['IDX'] = self.idx # the standalone glue token "++"
self.idx += 1
self.__build_vocab_lookup()
n_suffix = sum(1 for k in self.vocab_to_id
if k.startswith('++') and k != '++')
print(f" glue_morphemes: kept {n_suffix} learned ++suffix tokens, "
f"added one '++' glue token, skipped root-continuation duplicates "
f"(vocab = {self.get_vocab_size()}).")
def _reindex_contiguous(self) -> None:
"""Reassign token ids to a gap-free 0..N-1 range, keeping the five core
specials at their canonical ids. Vocab surgery (glue collapse) discards
ids and leaves holes; without this the model embedding (sized to
len(vocab)) could be indexed past its last row."""
self.__build_vocab_lookup()
core = ['<unk>', '<pad>', '<s>', '</s>', '<mask>']
old = self.vocab_to_id
new_id, nid = {}, 0
for t in core:
if t in old:
new_id[t] = nid; nid += 1
for t in sorted(old, key=lambda k: old[k]):
if t not in new_id:
new_id[t] = nid; nid += 1
def _remap(node, path):
for k, v in list(node.items()):
if k == 'IDX':
surf = ''.join(path).replace('[RSX]', '')
if surf in new_id:
node['IDX'] = new_id[surf]
elif isinstance(v, dict):
_remap(v, path + [k])
_remap(self.roots, [])
self.idx = nid
self.__build_vocab_lookup()
self.unk_token_id = new_id.get('<unk>', 0)
self.pad_token_id = new_id.get('<pad>', 1)
self.bos_token_id = new_id.get('<s>', 2)
self.eos_token_id = new_id.get('</s>', 3)
self.mask_token_id = new_id.get('<mask>', 4)
def __build_vocab_lookup(self) -> None:
self.vocab_to_id = {}
def traverse(trie, path):
for k, v in trie.items():
if k == 'IDX':
self.vocab_to_id[''.join(path)] = v
elif isinstance(v, dict):
traverse(v, path + [k])
traverse(self.roots, [])
self.vocab_to_id = dict(sorted(self.vocab_to_id.items(), key=lambda x: x[1]))
self.vocab_to_id = {k.replace("[RSX]", ""): v for k, v in self.vocab_to_id.items()}
self.id_to_vocab = {v: k for k, v in self.vocab_to_id.items()}
def __build_vocab_freq(self) -> None:
"""Build vocab_to_freq from trie nodes that have BOTH IDX and ## (in-training state)."""
vocab_to_freq = {}
def traverse(trie, path):
if isinstance(trie, dict) and 'IDX' in trie and '##' in trie:
vocab_to_freq[''.join(path)] = trie['##']
for k, v in trie.items():
if k not in ('IDX', '##') and isinstance(v, dict):
traverse(v, path + [k])
traverse(self.roots, [])
self.vocab_to_freq = dict(sorted(vocab_to_freq.items(), key=lambda x: x[1], reverse=True))
def __build_vocab_freq_snapshot(self) -> dict:
"""
Return a frequency-filtered vocabulary snapshot WITHOUT modifying self.roots.
Used exclusively by the type_based OOA loop. Unlike __build_vocab_freq
(which sets self.vocab_to_freq as a side-effect) this method returns a
new dict and leaves the trie completely unchanged — essential in
type_based mode where each word form is traversed only once and any
destructive cleaning would permanently corrupt the trie.
"""
snapshot = {}
def traverse(trie, path):
if isinstance(trie, dict) and 'IDX' in trie and '##' in trie:
freq = trie['##']
if freq > self.min_frequency:
snapshot[''.join(path)] = freq
for k, v in trie.items():
if k not in ('IDX', '##') and isinstance(v, dict):
traverse(v, path + [k])
traverse(self.roots, [])
return dict(sorted(snapshot.items(), key=lambda x: x[1], reverse=True))
def __sort_trie_by_freq(self, d):
if not isinstance(d, dict):
return d
stack = [d]
while stack:
node = stack.pop()
items = sorted(
node.items(),
key=lambda item: item[1].get('##', float('-inf')) if isinstance(item[1], dict) else float('-inf'),
reverse=True,
)
keys = list(node.keys())
for k in keys:
del node[k]
for k, v in items:
node[k] = v
if isinstance(v, dict):
stack.append(v)
return d
def __count_trie_nodes(self, trie_node) -> int:
"""Iterative (stack-based) node count. Must NOT be recursive: in
--boundaries_discovery mode the trie can have branches as deep as
the longest un-anchored sequence in the corpus, which routinely
exceeds Python's default recursion limit (RecursionError)."""
if not isinstance(trie_node, dict):
return 0
count = 0
stack = [trie_node]
while stack:
node = stack.pop()
count += 1
for key, value in node.items():
if key not in ('IDX', '##') and isinstance(value, dict):
stack.append(value)
return count
def prepare_encoding(self) -> None:
self.ids.append(0)
self.tokens.append("")
# =========================================================================
# Suffix-stem pruning
# =========================================================================
def __prune_weak_suffixes(self) -> None:
"""Two-pass ++ trie pruning (see v1.4.1 docstring for details)."""
if self.min_suffix_stems <= 0 or '++' not in self.roots:
return
pruned = 0
def _walk_prune(node: dict, path: str) -> None:
nonlocal pruned
for k in list(node.keys()):
if k in ('IDX', '##'):
continue
if isinstance(node[k], dict):
_walk_prune(node[k], path + k)
if 'IDX' in node:
n_stems = len(self.suffix_stems.get(path, set()))
if n_stems < self.min_suffix_stems:
del node['IDX']
pruned += 1
_walk_prune(self.roots['++'], '')
def _has_idx(node: dict) -> bool:
if 'IDX' in node:
return True
return any(
_has_idx(v) for k, v in node.items()
if isinstance(v, dict) and k not in ('IDX', '##')
)
for k in list(self.roots['++'].keys()):
if k in ('IDX', '##'):
continue
if isinstance(self.roots['++'][k], dict) and not _has_idx(self.roots['++'][k]):
del self.roots['++'][k]
if pruned:
print(f" Pruned {pruned} ++ suffix(es) with < {self.min_suffix_stems} distinct stems.")
# =========================================================================
# Training
# =========================================================================
# =========================================================================
# v1.4.5 — boundary discovery (whitespace-free, utterance-anchored)
# =========================================================================
def _anchor_sequences(self, text: str):
"""Cut `text` only at UNAMBIGUOUS edges: punctuation + line/tab breaks
(each is a true end "]]" of one sequence and a true start "[[" of the
next). Registered special tokens (e.g. CHILDES "*CHI:") are treated as
true starts. Inside a sequence, spaces are kept as the soft-cue symbol
SPACE_MARK rather than used as delimiters, so word boundaries must be
re-discovered from statistics (this is what makes the method work
identically for spaced eng/nld and unspaced zho)."""
pieces = []
split_re = self._special_token_splitter()
chunks = split_re.split(text) if split_re else [text]
for chunk in chunks:
if chunk in self.roots['[RSX]']:
continue # special = anchor only
for seg in self._anchor_re.split(chunk):
s = ''.join(ch for ch in seg
if ch.isalpha() or ch in (' ', "'", '-'))
s = re.sub(r'\s+', ' ', s).strip().replace(' ', self.SPACE_MARK)
if len(s) >= 2:
pieces.append(s)
return pieces
def __discover_boundaries(self, seq: str) -> None:
"""Single left-to-right pass over ONE true-anchored sequence. At each
position the root trie is read from the true start and the infl trie
from the true end; a morpheme break "||" is committed as soon as BOTH
respect the Tolerance Principle (sufficiency). Pre-req: `seq` has
already been built into self.roots (forward) and self.infls (reversed).
Mirrors the bookkeeping of __morsplit so emitted units are retrievable."""
n = len(seq)
if n < 3:
return
# v1.4.5a: O(n) instead of O(n^2). The split test at position i only
# needs consecutive node counts on the forward root path (freq(seq[:i-1]),
# freq(seq[:i])) and on the backward infl path (freq of the reversed
# suffixes of length n-i-1 and n-i). Both are obtained with ONE linear
# walk each, so we never re-slice/re-traverse from the anchors per i.
rc = [0] * (n + 1) # rc[k] = freq(seq[:k]) (root path)
rnode = [None] * (n + 1) # rnode[k]= node at seq[:k]
node = self.roots
for k in range(1, n + 1):
node = node[seq[k - 1]]
rc[k] = node['##']
rnode[k] = node
ic = [0] * (n + 1) # ic[j] = freq(seq[n-j:][::-1]) (infl path)
inode = [None] * (n + 1)
node = self.infls
rev = seq[::-1]
for j in range(1, n + 1):
node = node[rev[j - 1]]
ic[j] = node['##']
inode[j] = node
last = 0
for i in range(2, n): # commit "||" as soon as bilateral TP holds
if not self.__check_tp(rc[i - 1], rc[i]):
continue
if not self.__check_tp(ic[n - i - 1], ic[n - i]):
continue
morph, suffix = seq[last:i], seq[i:]
rn = rnode[i]
if 'IDX' not in rn:
rn['IDX'] = 1 # left segment retrievable
self.__build_trie(suffix, self.roots['++'])
self.suffix_stems.setdefault(suffix, set()).add(morph)
if self.ooa:
key = self.current_world + " " + morph + "||" + suffix
if key in self.ooa_split:
self.ooa_split[key] += 1
else:
self.ooa_data.append([
self.current_world, morph + "||" + suffix,
rc[i - 1], rc[i], len(rnode[i]),
ic[n - i - 1], ic[n - i], len(inode[n - i]),
len(self.types), self._ooa_exposure,
])
self.ooa_split[key] = 1
last = i
def __train_boundaries_discovery(self, all_contents: str, n_files: int,
output_dir: str) -> None:
"""Whitespace-free training pass: anchor -> sort shorter-first -> build
the dual trie over full sequences -> place boundaries by sufficiency ->
shared post-processing (identical to the tail of train())."""
sequences = self._anchor_sequences(all_contents)
sequences.sort(key=len) # developmental curriculum
self.num_tokens_in_corpus = len(sequences)
for s in sequences:
self.num_chars_in_corpus += len(s)
print(f"MorPiece (--boundaries_discovery): {n_files} files -> "
f"{len(sequences)} true-anchored sequences "
f"(whitespace NOT a boundary; shorter-first).")
token_counter = 0
next_ooa_at = self.ooa_token_interval
self._ooa_exposure = 0 # running sequence position for the OOA event log
for seq in sequences:
self.current_world = seq
self._ooa_exposure += 1
self.__build_trie(seq, self.roots) # root trie (true start)
self.__build_trie(seq[::-1], self.infls) # infl trie (true end)
self.__discover_boundaries(seq)
token_counter += self._ooa_increment(seq)
if self.ooa and token_counter >= next_ooa_at:
self.__incremental_cleaning(self.roots, self.min_frequency)
self.__build_vocab_freq()
ooa_dir = os.path.join(
output_dir, f"ooa_vocab_growth")
os.makedirs(ooa_dir, exist_ok=True)
self.save_vocab(os.path.join(ooa_dir, f"vocab_{token_counter}.json"))
print('.', end='', flush=True)
while next_ooa_at <= token_counter:
next_ooa_at += self.ooa_token_interval
if self.ooa:
print("")
# shared post-processing (mirror tail of train())
self.types = dict(sorted(self.types.items(), key=lambda x: x[1], reverse=True))
self.__sort_trie_by_freq(self.roots)
self.num_chars_in_trie = self.__count_trie_nodes(self.roots)
self.__prune_weak_suffixes()
self.__optimize(self.roots)
self.num_chars_in_optimized_trie = self.__count_trie_nodes(self.roots)
self._add_frequent_chars()
if self.glue_morphemes:
self._finalize_glue_vocab()
# else:
# self._add_root_continuations()
self._add_byte_fallback_tokens()
print(f"MorPiece (--boundaries_discovery) trained: final vocabulary = "
f"{self.get_vocab_size()} tokens")
# =========================================================================
# OOA exposure accounting (--counter_unit)
# =========================================================================
@staticmethod
def _count_syllables(text: str) -> int:
"""Approximate syllable count for OOA exposure accounting.
Latin script: number of maximal vowel runs (the standard cheap nucleus
estimate; slightly under-counts hiatus such as IT 'pa-u-ra'). CJK
ideographs, Kana morae and Hangul syllable blocks count 1 each (~one
syllable per glyph), keeping the unit meaningful for zho/jpn/kor under
--boundaries_discovery. This is for *setting an interval*, not for
linguistic analysis: the syllable is used because it is the most
cross-linguistically comparable measure of input quantity
(Raesaenen et al. 2019, 2021)."""
VOWELS = set("aeiouy\u00e0\u00e1\u00e2\u00e3\u00e4\u00e5\u00e8\u00e9"
"\u00ea\u00eb\u00ec\u00ed\u00ee\u00ef\u00f2\u00f3\u00f4"
"\u00f5\u00f6\u00f9\u00fa\u00fb\u00fc\u00fd\u00ff")
syl, in_vowel, has_letter = 0, False, False
for ch in text.lower():
cp = ord(ch)
if (0x4E00 <= cp <= 0x9FFF or 0x3400 <= cp <= 0x4DBF or
0xF900 <= cp <= 0xFAFF or 0x3040 <= cp <= 0x30FF or
0xAC00 <= cp <= 0xD7A3):
syl += 1
in_vowel = False
has_letter = True
continue
if ch in VOWELS:
has_letter = True
if not in_vowel:
syl += 1
in_vowel = True
else:
if ch.isalpha():
has_letter = True
in_vowel = False
if has_letter and syl == 0: # consonant-only alphabetic token
return 1
return syl
def _ooa_increment(self, text: str) -> int:
"""How much one processed item advances the token-based OOA exposure
counter, per self.counter_unit: 'tokens'->1, 'chars'->len, 'syllables'->~syl."""
if self.counter_unit == 'chars':
return len(text)
if self.counter_unit == 'syllables':
return self._count_syllables(text)
return 1
def train(self, training_corpus_path: str, text_column: str = "text",
save_complete_tries=None, output_dir: str = "tokenizer") -> None:
"""
Train the MorPiece tokenizer.
Parameters
----------
training_corpus_path : str
Path to a directory of UTF-8 text files OR a single .parquet file.
text_column : str
Column name in the parquet file (default 'text').
save_complete_tries : str | bool | None
If set, dump the COMPLETE, un-pruned root + inflection tries to a
JSON file *before* any min_frequency / min_suffix_stems pruning is
applied (see "Changes in 1.4.3"). Pass a path string, or `True`
for the default 'tokenizer/complete_tries.json'. Default None
(disabled) — training is unaffected when the option is not used.
The resulting file is consumed by morpiece_trie_explorer.py.
output_dir : str
Root directory for OOA snapshots (default 'tokenizer'). Both
type_based and token-based snapshots land in the same
'ooa_cutoff{c}_mfreq{m}' subfolder; the filename (vocab_type_* vs
vocab_*) carries the mode distinction.
"""
all_contents = ""
n_files = 0
training_corpus_path = str(training_corpus_path)
if training_corpus_path.endswith(".parquet"):
try:
import pandas as pd
except ImportError:
raise ImportError("pandas required: pip install pandas pyarrow")
df = pd.read_parquet(training_corpus_path)
if text_column not in df.columns:
raise ValueError(
f"Column '{text_column}' not found. Available: {df.columns.tolist()}"
)
for content in df[text_column].dropna().astype(str):
content = self._preprocess_text(content)
all_contents += content + "\n"
n_files += 1
print(
f"MorPiece tokenizer training: loaded {n_files} rows from "
f"parquet '{training_corpus_path}' (column='{text_column}')..."
)
else:
for filename in os.listdir(training_corpus_path):
file_path = os.path.join(training_corpus_path, filename)
if os.path.isfile(file_path):
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
content = self._preprocess_text(content)
all_contents += content + "\n"
n_files += 1
# v1.4.5: boundary-discovery mode never tokenises on whitespace.
if self.boundaries_discovery:
self.__train_boundaries_discovery(all_contents, n_files, output_dir)
return
words = all_contents.split()
print(
f"MorPiece tokenizer training: processing corpus "
f"({n_files} files, {len(words)} tokens)..."
)
if self.ooa:
if self.type_based:
print(
f"Saving Order Of Acquisition Vocabulary every "
f"{self.ooa_type_interval} types (frequency-ordered) (.):"
)
else:
print(f"Saving Order Of Acquisition Vocabulary every "
f"{self.ooa_token_interval} {self.counter_unit} (.):")
# --- type-frequency counting pass ------------------------------------
for word in words:
# v1.4.4: a registered special token (e.g. CHILDES "*CHI:") is
# atomic — count it as a corpus token but keep it out of `types`
# so it neither competes for vocab nor enters the morpheme trie.
if word in self.roots['[RSX]']:
self.num_tokens_in_corpus += 1
self.num_chars_in_corpus += len(word)
continue
word_alpha = ''.join(
ch for ch in word if ch.isalpha() or ch in ("'", "-")
)
word = word_alpha if word_alpha else word
if word:
self.types[word] = self.types.get(word, 0) + 1
self.num_tokens_in_corpus += 1
self.num_chars_in_corpus += len(word)
# --- choose training order -------------------------------------------
if self.type_based:
if self.ooa:
training_items = sorted(
self.types.keys(), key=lambda w: self.types[w], reverse=True
)
print(
f" type_based=True (ooa): iterating {len(training_items)} types "
f"in descending frequency order."
)
else:
training_items = list(self.types.keys())
print(
f" type_based=True: {len(training_items)} unique types "
f"(corpus {self.num_tokens_in_corpus} tokens, "
f"TTR={round(len(self.types)/self.num_tokens_in_corpus, 3)})"
)
else:
training_items = words
# --- main trie-building pass -----------------------------------------
token_counter = 0
next_ooa_at = self.ooa_token_interval
self._ooa_exposure = 0 # running token position for the OOA event log
for word in training_items:
# v1.4.4: registered special tokens stay atomic — they already
# live in [RSX] with their own IDX, so do not morsplit them.
if word in self.roots['[RSX]']:
self.current_world = word
self._ooa_exposure += 1
token_counter += 1 if self.type_based else self._ooa_increment(word)
continue
word_alpha = ''.join(
ch for ch in word if ch.isalpha() or ch in ("'", "-")
)
word = word_alpha if word_alpha else word
if not word:
continue
self.current_world = word
self._ooa_exposure += 1 # stamp this token's position before morsplit
self.__build_trie(word[::-1], self.infls)
self.__build_trie(word, self.roots)
self.__morsplit(word)
token_counter += 1 if self.type_based else self._ooa_increment(word)
# -----------------------------------------------------------------
# OOA snapshots
# type_based : every `ooa_type_interval` TYPES (exposure = types).
# token-based: every `ooa_token_interval` units of `counter_unit`
# (tokens | chars | syllables). Threshold-crossing
# (>=) instead of modulo, so chars/syllables — whose
# per-item increment exceeds 1 — never skip a snapshot.
# -----------------------------------------------------------------
if self.ooa:
ooa_dir = os.path.join(
output_dir, f"ooa_vocab_growth"
)
if self.type_based:
if token_counter % self.ooa_type_interval == 0:
snap = self.__build_vocab_freq_snapshot()
os.makedirs(ooa_dir, exist_ok=True)
snap_path = os.path.join(ooa_dir, f"vocab_type_{token_counter}.json")
with open(snap_path, 'w') as f:
json.dump({'vocab': snap}, f, indent=2)
print('.', end='', flush=True)
elif token_counter >= next_ooa_at:
self.__incremental_cleaning(self.roots, self.min_frequency)
self.__build_vocab_freq()
os.makedirs(ooa_dir, exist_ok=True)
self.save_vocab(os.path.join(ooa_dir, f"vocab_{token_counter}.json"))
print('.', end='', flush=True)
while next_ooa_at <= token_counter:
next_ooa_at += self.ooa_token_interval
if self.ooa:
print("")
self.types = dict(sorted(self.types.items(), key=lambda x: x[1], reverse=True))
self.__sort_trie_by_freq(self.roots)
self.num_chars_in_trie = self.__count_trie_nodes(self.roots)
# --- v1.4.3: optional snapshot of the COMPLETE (un-pruned) tries -----
# Taken here, after frequency-sorting but BEFORE __prune_weak_suffixes()
# and BEFORE __optimize() — i.e. before any min_suffix_stems / min_freq
# pruning. Every node still carries its raw '##' frequency count.
if save_complete_tries:
ct_path = (save_complete_tries if isinstance(save_complete_tries, str)
else os.path.join(output_dir, "complete_tries.json"))
self.save_complete_tries(ct_path)
self.__prune_weak_suffixes()
print(
f"MorPiece tokenizer training: trie optimisation "
f"(vocab_size={self.vocab_size}, min_freq={self.min_frequency}, "
f"cutoff={self.cutoff}, "
f"min_suffix_stems={self.min_suffix_stems})..."
)
self.__optimize(self.roots)
self.num_chars_in_optimized_trie = self.__count_trie_nodes(self.roots)
self._add_frequent_chars()
if self.glue_morphemes:
self._finalize_glue_vocab()
else:
self._add_root_continuations()
self._add_byte_fallback_tokens()
print(f"MorPiece tokenizer trained: final vocabulary = {self.get_vocab_size()} tokens")
if self.ooa:
print("MorPiece tokenizer Order of Acquisition data prepared.")
# =========================================================================
# Encoding / Decoding
# =========================================================================
def encode(self, sentence: str):
"""Encode a sentence into (token_ids, token_strings)."""
sentence = self._preprocess_text(sentence)
self.ids, self.tokens = [], []
for word in sentence.strip().split():
if word in self.roots['[RSX]']:
self.ids.append(self.roots['[RSX]'][word]['IDX'])
self.tokens.append(word)
else:
self.prepare_encoding()
self.__retrieve(word, self.roots)
if self.use_tokenizers_lib:
n_before = len(self.ids)
self.ids = self._postprocess_tokens(self.ids)
if len(self.ids) == n_before + 1:
self.tokens = [self.start_of_text_symbol] + self.tokens
return self.ids, self.tokens
def decode(self, sentence_idxs: list) -> list:
raw = [self.id_to_vocab.get(idx, '<unk>') for idx in sentence_idxs]
tokens = self._fuse_byte_tokens(raw) if self.byte_fallback else raw
if self.use_tokenizers_lib and tokens:
text = ''.join(tokens)
try:
return self.decoder.decode(text).split()
except Exception:
return tokens
return tokens
def _fuse_byte_tokens(self, tokens: list) -> list:
"""Collapse runs of "<0xHH>" byte tokens back into UTF-8 text."""
out, buf = [], []
for t in tokens:
if t in self._byte_tok_set:
buf.append(int(t[3:5], 16))
else:
if buf:
out.append(bytes(buf).decode('utf-8', errors='replace'))
buf = []
out.append(t)
if buf:
out.append(bytes(buf).decode('utf-8', errors='replace'))
return out
# =========================================================================
# Special-token management
# =========================================================================
def set_special_tokens(self, token_list: list) -> None:
for item in token_list:
if item not in self.roots['[RSX]']:
self.roots['[RSX]'][item] = {'IDX': self.idx}
self.idx += 1
def pad_sentence(self, sentence: str, l: int, pad: str = '<pad>') -> str:
words = sentence.split()
n_pad = max(l - len(words), 0)
return ' '.join([pad] * n_pad + words)
# =========================================================================
# Diagnostics
# =========================================================================
def diagnose_tp(self, sample_words: list = None, n_words: int = 20) -> None:
"""
Print a per-word trace of why TP does or does not fire, to help
calibrate cutoff / type_based / min_suffix_stems parameters.
Uses the suffix_stems dictionary (populated during training) together
with the vocabulary to report whether splits were found and how many
stems each suffix derives from. Also calls encode() on each sample
word to show the final tokenization.
Parameters
----------
sample_words : list of str, optional
Words to trace. If None, the 20 most-frequent types are used.
n_words : int
How many words to show when sample_words is None (default 20).
Example
-------
>>> mop.diagnose_tp(['cantava', 'telefono', 'asociale', 'asola'])
"""
if not self.types:
print("No type counts available — train first.")
return
if sample_words is None:
sample_words = [w for w, _ in
sorted(self.types.items(), key=lambda x: -x[1])[:n_words]]
print(f"diagnose_tp cutoff={self.cutoff} type_based={self.type_based} "
f"min_suffix_stems={self.min_suffix_stems}\n")
stem_suffix_pairs = set()
for suf, stems in self.suffix_stems.items():
for s in stems:
stem_suffix_pairs.add((s, suf))
for w in sample_words:
ids_out, toks_out = self.encode(w)
in_vocab = (w in (self.vocab_to_id or {})) or (w in self.roots.get('[RSX]', {}))
freq = self.types.get(w, 0)
print(f" Word: {repr(w)} (corpus freq={freq})")
print(f" encode → {toks_out}")
print(f" in vocabulary: {in_vocab}")
found_splits = [(s, suf) for s, suf in stem_suffix_pairs
if s + suf == w and len(s) >= 1 and len(suf) >= 1]
if found_splits:
print(f" Splits proposed during training:")
for stem, suf in sorted(found_splits, key=lambda x: len(x[0])):
n_stems = len(self.suffix_stems.get(suf, set()))
survives = n_stems >= self.min_suffix_stems
status = (f"++ kept ({n_stems} stems ≥ {self.min_suffix_stems})"
if survives else
f"++ PRUNED ({n_stems} stems < {self.min_suffix_stems})")
print(f" {repr(stem)} | {repr(suf)} → {status}")
else:
reasons = []
if freq == 0:
reasons.append("word never seen in corpus")
else:
reasons.append("no split fired during training")
reasons.append(f"possible causes: cutoff too high (need freq>{self.cutoff}); "
f"word too short or non-decomposable; "
f"infl-trie ending below cutoff")
print(f" No splits proposed: {'; '.join(reasons)}")
print()
print(f" === Global summary ===")
print(f" Total ++ suffixes proposed: {len(self.suffix_stems)}")
top = sorted(self.suffix_stems.items(), key=lambda x: -len(x[1]))[:10]
for suf, stems in top:
survives = len(stems) >= self.min_suffix_stems
mark = "✓" if survives else "✗"
print(f" {mark} {repr(suf):12s} {len(stems):3d} stems e.g. {sorted(list(stems))[:3]}")
n_pp = len([t for t in (self.vocab_to_id or {}) if t.startswith('++')])
print(f" ++ tokens in final vocabulary: {n_pp}")
print(f" Parameter advice:")
if n_pp == 0 and len(self.suffix_stems) == 0:
print(f" No splits fired at all.")
print(f" → Lower cutoff (current={self.cutoff}) if corpus is small.")
print(f" → Ensure type_based=False for token-frequency-driven TP.")
elif n_pp == 0:
print(f" Splits proposed but all pruned by min_suffix_stems={self.min_suffix_stems}.")
print(f" → Lower min_suffix_stems or expand corpus for more stem diversity.")
else:
print(f" ++ tokens present — tokenizer is segmenting. Tune to taste.")
# =========================================================================
# Statistics
# =========================================================================
def get_num_chars_in_trie(self): return self.num_chars_in_trie
def get_num_chars_in_optimized_trie(self): return self.num_chars_in_optimized_trie
def get_num_chars_in_corpus(self): return self.num_chars_in_corpus
def get_num_tokens_in_corpus(self): return self.num_tokens_in_corpus
def get_num_types_in_corpus(self): return len(self.types)
def get_compression_ratio(self): return round(self.num_chars_in_optimized_trie / self.num_chars_in_trie * 100, 3)
def get_ttr(self): return round(len(self.types) / self.num_tokens_in_corpus, 3)
def get_vocab_size(self) -> int:
self.vocab_size = self.idx
return self.idx
# =========================================================================
# Serialisation
# =========================================================================
def from_pretrained(self, load_file: str) -> None:
tokenizer_path = os.path.join(load_file, 'tokenizer.json')
print(f"Loading MorPiece tokenizer from {tokenizer_path}...")
with open(tokenizer_path, 'r', encoding="utf-8") as f:
data = json.load(f)
if isinstance(data, dict) and 'roots' in data:
self.roots = data['roots']
self.vocab_to_id = data.get('vocab', {})
self.id_to_vocab = {v: k for k, v in self.vocab_to_id.items()}
if self.vocab_to_id:
self.idx = max(self.vocab_to_id.values()) + 1
if 'special_token_ids' in data:
sp = data['special_token_ids']
self.unk_token_id = sp.get('unk', self.unk_token_id)
self.pad_token_id = sp.get('pad', self.pad_token_id)
self.bos_token_id = sp.get('bos', self.bos_token_id)
self.eos_token_id = sp.get('eos', self.eos_token_id)
self.mask_token_id = sp.get('mask', self.mask_token_id)
self.glue_morphemes = bool(
data.get('glue_morphemes', '++' in self.vocab_to_id))
self.glue_cjk_prefer_root = bool(data.get('glue_cjk_prefer_root', True))
else:
self.roots = data
self.vocab_to_id = {}
self.id_to_vocab = {}
if '[RSX]' not in self.roots:
raise ValueError("Invalid tokenizer format: missing [RSX] root node.")
# A loaded tokenizer may carry special tokens that the splitter cache
# has not seen — force a rebuild on next _preprocess_text call.
self._split_re_n = -1
def save_pretrained(self, save_file: str) -> None:
self.__build_vocab_lookup()
with open(save_file, 'w', encoding='utf-8') as f:
json.dump({
'roots': self.roots,
'vocab': self.vocab_to_id,
'glue_morphemes': bool(getattr(self, 'glue_morphemes', False)),
'glue_cjk_prefer_root': bool(getattr(self, 'glue_cjk_prefer_root', True)),
'special_token_ids': {
'unk': self.unk_token_id, 'pad': self.pad_token_id,
'bos': self.bos_token_id, 'eos': self.eos_token_id,
'mask': self.mask_token_id,
},
}, f, indent=2)
def save_complete_tries(self, save_file: str) -> None:
"""
Dump the COMPLETE, UN-PRUNED tries to a single JSON file (v1.4.3).
Called from train() before __prune_weak_suffixes() and __optimize()
when the save_complete_tries option is requested. The dump preserves:
• every node's raw '##' frequency count;
• the placeholder 'IDX' markers (a registered word/morpheme path);
• the reversed-word inflection trie (self.infls);
• the ++ suffix→stems map backing min_suffix_stems pruning.
Together these let a downstream tool reconstruct exactly which
branches the later min_frequency cut will remove and what frequency
each surviving / discarded node carried. Consumed by
morpiece_trie_explorer.py.
NOTE: in token-based mode (type_based=False) with ooa=True the trie
has already been thinned every 100 000 tokens by
__incremental_cleaning, so 'complete' there means 'complete as of
end-of-training', not 'every node ever created'.
"""
out_dir = os.path.dirname(save_file)
if out_dir:
os.makedirs(out_dir, exist_ok=True)
payload = {
"format": "morpiece-complete-tries/1",
"meta": {
"version": __version__,
"vocab_size": self.vocab_size,
"min_frequency": self.min_frequency,
"cutoff": self.cutoff,
"min_suffix_stems": self.min_suffix_stems,
"type_based": self.type_based,
"num_tokens_in_corpus": self.num_tokens_in_corpus,
"num_types_in_corpus": len(self.types),
},
"roots": self.roots,
"infls": self.infls,
"suffix_stems": {suf: sorted(stems)
for suf, stems in self.suffix_stems.items()},
}
with open(save_file, 'w', encoding='utf-8') as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
print(f"MorPiece complete (un-pruned) tries saved → {save_file}")
print(f" root nodes={self.__count_trie_nodes(self.roots)} "
f"infl nodes={self.__count_trie_nodes(self.infls)} "
f"++ suffixes={len(self.suffix_stems)}")
def save_vocab(self, save_file: str) -> None:
self.__sort_trie_by_freq(self.roots)
self.__build_vocab_lookup()
with open(save_file, 'w') as f:
json.dump({'vocab': self.vocab_to_freq}, f, indent=2)
def save_types(self, file: str) -> None:
with open(file, 'w') as f:
json.dump(self.types, f, indent=2)
def save_ooa(self, file: str) -> None:
with open(file, 'w', encoding='utf8') as f:
for row in self.ooa_data:
f.write('\t'.join(str(cell) for cell in row) + '\n')
def save_HF(self, save_directory: str, model_max_length: int = 1024) -> None:
"""Save in HuggingFace PreTrainedTokenizerFast format."""
os.makedirs(save_directory, exist_ok=True)
if not hasattr(self, 'vocab_to_id') or self.vocab_to_id is None:
self.__build_vocab_lookup()
vocab = [""] * len(self.vocab_to_id)
for token, token_id in self.vocab_to_id.items():
if token_id < len(vocab):
vocab[token_id] = token
for i in range(len(vocab)):
if vocab[i] == "":
vocab[i] = "<unk>"
sp = {
"<unk>": self.unk_token_id, "<pad>": self.pad_token_id,
"<s>": self.bos_token_id, "</s>": self.eos_token_id,
"<mask>": self.mask_token_id,
}
for token, token_id in sp.items():
if token_id < len(vocab):
vocab[token_id] = token
tokenizer_json = {
"version": "1.0", "truncation": None, "padding": None,
"added_tokens": [
{"id": sp[tok], "content": tok, "single_word": False,
"lstrip": False, "rstrip": False, "normalized": False, "special": True}
for tok in ("<unk>", "<pad>", "<s>", "</s>", "<mask>")
],
"normalizer": {
"type": "Sequence",
"normalizers": [{"type": "Lowercase"}, {"type": "NFKC"}],
},
"pre_tokenizer": {"type": "Whitespace"},
"post_processor": {
"type": "TemplateProcessing",
"single": [
{"SpecialToken": {"id": "<s>", "type_id": 0}},
{"Sequence": {"id": "A", "type_id": 0}},
],
"pair": [
{"SpecialToken": {"id": "<s>", "type_id": 0}},
{"Sequence": {"id": "A", "type_id": 0}},
{"SpecialToken": {"id": "<s>", "type_id": 1}},
{"Sequence": {"id": "B", "type_id": 1}},
],
"special_tokens": {
"<s>": {"id": "<s>", "ids": [sp["<s>"]], "tokens": ["<s>"]},
"</s>": {"id": "</s>", "ids": [sp["</s>"]], "tokens": ["</s>"]},
},
},
"decoder": {"type": "WordPiece", "prefix": "++", "cleanup": True},
"model": {
"type": "WordPiece", "unk_token": "<unk>",
"continuing_subword_prefix": "++",
"max_input_chars_per_word": 100,
"vocab": {token: i for i, token in enumerate(vocab) if token.strip()},
},
}
with open(os.path.join(save_directory, "tokenizer.json"), 'w', encoding='utf-8') as f:
json.dump(tokenizer_json, f, indent=2, ensure_ascii=False)
added_tokens_decoder = {
str(tid): {"content": tok, "lstrip": False, "normalized": False,
"rstrip": False, "single_word": False, "special": True}
for tok, tid in sp.items()
}
tokenizer_config = {
"added_tokens_decoder": added_tokens_decoder,
"bos_token": "<s>", "clean_up_tokenization_spaces": False,
"eos_token": "</s>", "extra_special_tokens": {},
"mask_token": "<mask>", "model_max_length": model_max_length,
"pad_token": "<pad>", "tokenizer_class": "PreTrainedTokenizerFast",
"unk_token": "<unk>",
}
with open(os.path.join(save_directory, "tokenizer_config.json"), 'w', encoding='utf-8') as f:
json.dump(tokenizer_config, f, indent=2, ensure_ascii=False)
special_tokens_map = {
"bos_token": "<s>", "eos_token": "</s>",
"unk_token": "<unk>", "pad_token": "<pad>", "mask_token": "<mask>",
}
with open(os.path.join(save_directory, "special_tokens_map.json"), 'w', encoding='utf-8') as f:
json.dump(special_tokens_map, f, indent=2, ensure_ascii=False)
with open(os.path.join(save_directory, "vocab.txt"), 'w', encoding='utf-8') as f:
for token in vocab:
if token.strip():
f.write(token + "\n")
print(f"MorPiece tokenizer saved in HuggingFace format → {save_directory}")
print(f" vocab_size={len([t for t in vocab if t.strip()])} model_max_length={model_max_length}")
print(f" Load: PreTrainedTokenizerFast.from_pretrained('{save_directory}')")
# =========================================================================
# HF Tokenizer factory helpers
# =========================================================================
def create_WordPiece_tokenizer(self):
if not self.use_tokenizers_lib:
print("Tokenizers library integration is disabled.")
return None
from tokenizers import Tokenizer
from tokenizers.models import WordPiece
if not hasattr(self, 'vocab_to_id') or self.vocab_to_id is None:
self.__build_vocab_lookup()
vocab = [""] * len(self.vocab_to_id)
for token, token_id in self.vocab_to_id.items():
if token_id < len(vocab):
vocab[token_id] = token
for i in range(len(vocab)):
if vocab[i] == "":
vocab[i] = "<unk>"
tokenizer = Tokenizer(WordPiece({t: i for i, t in enumerate(vocab)}, unk_token="<unk>"))
tokenizer.normalizer = self.normalizer
tokenizer.pre_tokenizer = self.pre_tokenizer
tokenizer.decoder = self.decoder
tokenizer.post_processor = self.post_processor
tokenizer.save()
return tokenizer
def create_unigram_tokenizer(self, verbose: bool = True, cluster_power: float = 2.0):
"""Export the trained vocab as a *fast* HF Unigram tokenizer.
Unigram runs Viterbi over per-token scores, so (unlike BPE) it needs no
merge chains and can use ANY subset of the vocab. It loads with a bare
AutoTokenizer.from_pretrained(dir) — no trust_remote_code, no custom .py.
MAXIMISING token clustering within a word: Viterbi maximises the SUM of
piece scores, and any score that is linear in length sums to the (fixed)
word length for every partition — so it cannot prefer one segmentation
over another. We therefore score each token SUPERLINEARLY in its length,
score = len**cluster_power (+ a small frequency tie-break). With
cluster_power>1 a single long piece beats the sum of its parts
(5**2=25 > 3**2+2**2=13), so Viterbi concentrates characters into the
longest available morpheme. Byte tokens get a large negative score
(last resort only).
Caveat: a stock model matches token SURFACE strings against the text, and
MorPiece "++X" continuation tokens never appear literally in the surface,
so they are exported by their SURFACE form ("++ed" -> "ed") and the "++"
glue is dropped. The inflection-vs-root distinction therefore exists only
in the native tokenizer, not in this export.
"""
try:
from tokenizers import Tokenizer, decoders as _dec
from tokenizers.models import Unigram
except ImportError:
print("tokenizers library not available."); return None
if not getattr(self, 'vocab_to_id', None):
self.__build_vocab_lookup()
import math
freq = self.vocab_to_freq or {}
specials = ['<unk>', '<pad>', '<s>', '</s>', '<mask>']
# Map every non-special, non-byte token to its SURFACE form, keeping the
# highest frequency seen for that surface (root and "++surface" collapse).
surface_freq = {}
for tok in self.vocab_to_id:
if tok in specials or tok == '++' or tok in self.reserved_keys \
or tok in self.BYTE_TOKENS:
continue
surf = tok[2:] if tok.startswith('++') else tok
if surf:
surface_freq[surf] = max(surface_freq.get(surf, 0), freq.get(tok, 0))
entries = [(s, 0.0) for s in specials] # unk_id = 0
for surf, f in surface_freq.items():
entries.append((surf, float(len(surf) ** cluster_power
+ 0.25 * math.log1p(f))))
for bt in self.BYTE_TOKENS: # last resort
entries.append((bt, -1.0e4))
uni = Unigram(entries, unk_id=0, byte_fallback=True)
tok = Tokenizer(uni)
tok.normalizer = self.normalizer
tok.pre_tokenizer = self.pre_tokenizer
tok.post_processor = self.post_processor
tok.decoder = _dec.Sequence([_dec.ByteFallback(), _dec.Fuse()])
if verbose:
print(f"Unigram tokenizer: {len(entries)} pieces (surface forms, "
f"cluster_power={cluster_power}, byte_fallback).")
return tok
def save_unigram_pretrained(self, save_directory: str,
model_max_length: int = 1024, cluster_power: float = 2.0):
"""Build the Unigram tokenizer and save a directory loadable with a plain
AutoTokenizer.from_pretrained(save_directory) (no trust_remote_code)."""
import os, json
tok = self.create_unigram_tokenizer(cluster_power=cluster_power)
if tok is None:
return None
os.makedirs(save_directory, exist_ok=True)
tok.save(os.path.join(save_directory, "tokenizer.json"))
with open(os.path.join(save_directory, "tokenizer_config.json"), "w",
encoding="utf-8") as f:
json.dump({"tokenizer_class": "PreTrainedTokenizerFast",
"model_max_length": model_max_length,
"bos_token": "<s>", "eos_token": "</s>", "unk_token": "<unk>",
"pad_token": "<pad>", "mask_token": "<mask>",
"clean_up_tokenization_spaces": False}, f, indent=2)
with open(os.path.join(save_directory, "special_tokens_map.json"), "w",
encoding="utf-8") as f:
json.dump({"bos_token": "<s>", "eos_token": "</s>", "unk_token": "<unk>",
"pad_token": "<pad>", "mask_token": "<mask>"}, f, indent=2)
print(f" saved Unigram tokenizer -> {save_directory} "
f"(load: AutoTokenizer.from_pretrained('{save_directory}'))")
return save_directory
def unigram_clustering(self, words, save_directory=None, cluster_power: float = 2.0):
"""Report clustering (pieces/word, chars/piece) for the Unigram export and
compare to native. Higher chars/piece = more clustering."""
from tokenizers import Tokenizer
import os
if save_directory:
ut = Tokenizer.from_file(os.path.join(save_directory, "tokenizer.json"))
else:
ut = self.create_unigram_tokenizer(verbose=False, cluster_power=cluster_power)
u_pieces = n_native = same_count = tot_chars = 0
for w in words:
up = [t for t in ut.encode(w, add_special_tokens=False).tokens]
_, nt = self.encode(w)
nt = [t for t in nt if t != self.start_of_text_symbol]
u_pieces += len(up); n_native += len(nt); tot_chars += len(w)
if len(up) <= len(nt):
same_count += 1
print(f"Unigram clustering over {len(words)} words: "
f"{u_pieces/len(words):.2f} pieces/word "
f"({tot_chars/max(1,u_pieces):.2f} chars/piece), "
f"native {n_native/len(words):.2f} pieces/word; "
f"unigram <= native on {same_count}/{len(words)}.")
return u_pieces / max(1, len(words)) |