Upload trained ModernMolBERT checkpoint

README.md

@@ -1,53 +1,114 @@
 ---
+license: mit
 library_name: transformers
+pipeline_tag: fill-mask
 tags:
-  - chemistry
-  - selfies
-  - masked-language-modeling
+- chemistry
+- molecules
+- selfies
+- ape-tokenizer
+- modernbert
+- masked-language-modeling
 ---
 
 # HauserGroup/ModernMolBERT-small
 
-ModernBERT pre-trained on SELFIES for masked language modeling.
+ModernMolBERT is a compact ModernBERT encoder pre-trained from scratch with masked language modeling on ~2.4M SELFIES strings from ChEMBL 36, using a chemically aware Atom Pair Encoding (APE) tokenizer. It expects SELFIES input and produces general-purpose molecular embeddings.
 
-## Model details
+## Model Details
+
+- **Developed by:** Hauser Group, Department of Drug Design and Pharmacology, University of Copenhagen
+- **Model type:** ModernBERT encoder — molecular embedding model trained with masked language modeling
+- **Input representation:** SELFIES (convert SMILES first; see below)
+- **Tokenizer:** Atom Pair Encoding (APE) over SELFIES primitives
+- **Pre-training data:** ChEMBL 36 (~2.4M unique small molecules)
+- **License:** MIT
+- **Repository:** https://github.com/HauserGroup/ModernMolBERT
 
 | field | value |
 |-------|-------|
 | model_type | modernbert |
-| vocab_size | 2439 |
+| vocab_size | 631 |
 | hidden_size | 512 |
 | num_hidden_layers | 8 |
 | num_attention_heads | 8 |
 | intermediate_size | 2048 |
 | max_position_embeddings | 128 |
 
-## Training
+## How to Get Started with the Model
 
-| field | value |
-|-------|-------|
-| dataset_name | data/pretrain/chembl36_selfies_tokenized |
-| model_size | small |
-| mlm_probability | 0.15 |
-| masking_strategy | standard |
-| max_steps | 30000 |
-| learning_rate | 0.0001 |
-| per_device_train_batch_size | 256 |
-| gradient_accumulation_steps | 1 |
-| warmup_steps | 1500 |
-| weight_decay | 0.01 |
-| max_seq_length | 128 |
-| seed | 42 |
-| best_eval_loss | 1.500394 |
-| best_global_step | 25000 |
-
-## Usage
+The model consumes **SELFIES** strings tokenized with the APE tokenizer. The main output for molecular representation learning is the first-token embedding:
+
+```python
+# pip install transformers torch
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+repo = 'HauserGroup/ModernMolBERT-small'
+model = AutoModel.from_pretrained(repo).eval()
+tokenizer = AutoTokenizer.from_pretrained(
+    repo,
+    subfolder='ape_tokenizer',
+    trust_remote_code=True,
+    use_fast=False,
+)
+
+# A SELFIES string (one bracketed token per primitive); here psilocybin.
+selfies = '[C][N][Branch1][C][C][C][C][C][=C][NH1][C][=C][C][=C][C][Branch1][#Branch2][O][P][=Branch1][C][=O][Branch1][C][O][O][=C][Ring1][=C][Ring1][O]'
+
+inputs = tokenizer(selfies, return_tensors='pt')
+with torch.no_grad():
+    outputs = model(**inputs)
+    embedding = outputs.last_hidden_state[:, 0]
+
+tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
+embedding_preview = [round(x, 4) for x in embedding[0, :5].tolist()]
+print(f"Token IDs:\n{inputs['input_ids'][0].tolist()}\n")
+print(f"Tokens:\n{tokens}\n")
+print(f"Embedding shape: {tuple(embedding.shape)}")
+print(f"Embedding first 5 values:\n{embedding_preview}")
+```
+
+Output:
+
+```text
+Token IDs:
+[0, 352, 336, 334, 334, 7, 406, 388, 388, 392, 489, 335, 18, 336, 426, 482, 482, 6, 2]
+
+Tokens:
+['<s>', '[C][N]', '[Branch1][C]', '[C][C]', '[C][C]', '[=C]', '[NH1][C]', '[=C][C]', '[=C][C]', '[Branch1][#Branch2]', '[O][P]', '[=Branch1][C]', '[=O]', '[Branch1][C]', '[O][O]', '[=C][Ring1]', '[=C][Ring1]', '[O]', '</s>']
+
+Embedding shape: (1, 512)
+Embedding first 5 values:
+[-0.1029, 0.2197, -0.0518, -0.7983, -0.6783]
+```
+
+If you start from SMILES, convert it to SELFIES first (e.g. the [`selfies`](https://github.com/aspuru-guzik-group/selfies) package: `selfies.encoder("CC(=O)Oc1ccccc1C(=O)O")`).
+
+For masked-token predictions, load the same checkpoint with `AutoModelForMaskedLM`:
 
 ```python
-from transformers import AutoModelForMaskedLM, AutoTokenizer
+from transformers import AutoModelForMaskedLM
+
+mlm = AutoModelForMaskedLM.from_pretrained(repo)
+logits = mlm(**inputs).logits
+print(f"Logits shape: {tuple(logits.shape)}")
+```
+
+Output:
 
-model = AutoModelForMaskedLM.from_pretrained("HauserGroup/ModernMolBERT-small")
-tokenizer = AutoTokenizer.from_pretrained("HauserGroup/ModernMolBERT-small", trust_remote_code=True)
+```text
+Logits shape: (1, 19, 631)
 ```
 
-This model expects SELFIES strings. Convert SMILES before tokenization.
+> Current Transformers releases disable custom root tokenizers for `model_type='modernbert'` before loading `auto_map`, so the tokenizer must be loaded from `ape_tokenizer/`. The root tokenizer files are also shipped for forward compatibility.
+
+## Uses
+
+- **Direct use:** molecular embeddings for property prediction, similarity search, clustering, and retrieval; masked-token fill-in.
+- **Downstream use:** fine-tuning for molecular classification or regression on SELFIES inputs.
+- **Out of scope:** natural-language text; generating valid SMILES; 3D/conformer-dependent tasks.
+
+## Bias, Risks, and Limitations
+
+Pre-trained only on drug-like ChEMBL 36 chemistry; may not generalize to natural products, agrochemicals, fragments, or other under-represented chemical space. Performance depends on the downstream task and adaptation strategy. No access to 3D/conformer information.

all_results.json

@@ -0,0 +1,16 @@
+{
+    "epoch": 3.066,
+    "eval_loss": 0.3744058609008789,
+    "eval_masked_accuracy": 0.8805643239064613,
+    "eval_perplexity": 1.454127204085947,
+    "eval_runtime": 2.8656,
+    "eval_samples_per_second": 1429.379,
+    "eval_steps_per_second": 5.584,
+    "num_parameters": 34149495.0,
+    "total_flos": 9.74439612930639e+16,
+    "train_loss": 0.6228721523284912,
+    "train_runtime": 2832.8619,
+    "train_samples_per_second": 2711.039,
+    "train_samples_streaming": 7680000.0,
+    "train_steps_per_second": 10.59
+}

ape_tokenizer/ape_tokenizer_metadata.json

@@ -0,0 +1,102 @@
+{
+  "dataset_name": "data/pretrain/chembl36_selfies",
+  "selfies_column": "selfies",
+  "train_split": "train",
+  "validation_split": "valid",
+  "use_validation_split": true,
+  "representation": "SELFIES",
+  "expected_input": "SELFIES strings only. Convert SMILES before inference using a helper such as smiles_to_selfies().",
+  "tokenizer_vocab_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+  "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json",
+  "backend": "cuda",
+  "platform": "Linux-6.8.0-110-generic-x86_64-with-glibc2.35",
+  "torch_version": "2.11.0+cu130",
+  "transformers_version": "5.8.1",
+  "vocab_size": 631,
+  "special_ids": {
+    "pad_token": 1,
+    "bos_token": 0,
+    "eos_token": 2,
+    "unk_token": 3,
+    "mask_token": 4
+  },
+  "num_parameters": 34149495,
+  "tokenizer_stats": {
+    "sample_size": 1000.0,
+    "mean_len": 25.555,
+    "p50_len": 25.0,
+    "p95_len": 41.0,
+    "p99_len": 51.0,
+    "max_len": 62.0,
+    "truncation_rate": 0.0,
+    "unk_rate": 0.0,
+    "empty_sequence_rate": 0.0,
+    "mostly_unknown_rate": 0.0
+  },
+  "final_eval_metrics": {
+    "eval_loss": 0.3744058609008789,
+    "eval_masked_accuracy": 0.8805643239064613,
+    "eval_runtime": 2.8656,
+    "eval_samples_per_second": 1429.379,
+    "eval_steps_per_second": 5.584,
+    "epoch": 3.066,
+    "eval_perplexity": 1.454127204085947
+  },
+  "trainer_state_summary": {
+    "best_global_step": 30000,
+    "best_metric": 0.37694016098976135,
+    "best_model_checkpoint": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_4e-4/checkpoint-30000",
+    "global_step": 30000
+  },
+  "args": {
+    "output_dir": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_4e-4",
+    "tokenizer_vocab_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+    "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json",
+    "dataset_name": "data/pretrain/chembl36_selfies",
+    "selfies_column": "selfies",
+    "train_split": "train",
+    "validation_split": "valid",
+    "use_validation_split": true,
+    "data_dir": null,
+    "data_files": null,
+    "eval_size": 4096,
+    "shuffle_buffer_size": 100000,
+    "seed": 42,
+    "val_split_mod": 100,
+    "val_split_bucket": 0,
+    "tokenizer_validation_samples": 1000,
+    "unk_rate_threshold": 0.001,
+    "truncation_warn_threshold": 0.05,
+    "model_size": "small",
+    "max_seq_length": 128,
+    "mlm_probability": 0.15,
+    "masking_strategy": "standard",
+    "span_p_geom": 0.4,
+    "span_max_length": 6,
+    "heteroatom_start_weight": 2.0,
+    "max_steps": 30000,
+    "per_device_train_batch_size": 256,
+    "per_device_eval_batch_size": 256,
+    "gradient_accumulation_steps": 1,
+    "learning_rate": 0.0004,
+    "weight_decay": 0.01,
+    "warmup_steps": 1500,
+    "max_grad_norm": 1.0,
+    "load_best_model_at_end": true,
+    "metric_for_best_model": "eval_loss",
+    "greater_is_better": false,
+    "logging_steps": 100,
+    "eval_steps": 5000,
+    "save_steps": 5000,
+    "save_total_limit": 2,
+    "device_backend": "cuda",
+    "bf16": true,
+    "fp16": false,
+    "num_workers": 4,
+    "max_eval_batches": 16,
+    "report_to": "tensorboard",
+    "compute_masked_accuracy": true,
+    "debug": false,
+    "hf_login": false
+  }
+}

ape_tokenizer/selfies_vocab.json

@@ -0,0 +1,633 @@
+{
+    "</s>": 2,
+    "<mask>": 4,
+    "<pad>": 1,
+    "<s>": 0,
+    "<unk>": 3,
+    "[#11C]": 256,
+    "[#14C]": 247,
+    "[#Branch1]": 22,
+    "[#Branch1][=C]": 479,
+    "[#Branch1][=O]": 542,
+    "[#Branch1][C]": 400,
+    "[#Branch1][N]": 453,
+    "[#Branch1][O]": 491,
+    "[#Branch1][S]": 528,
+    "[#Branch2]": 15,
+    "[#Branch2][=C]": 469,
+    "[#Branch2][=N]": 561,
+    "[#Branch2][=O]": 531,
+    "[#Branch2][C]": 380,
+    "[#Branch2][N]": 438,
+    "[#Branch2][O]": 461,
+    "[#Branch2][S]": 529,
+    "[#C-1]": 141,
+    "[#CH0]": 248,
+    "[#C]": 23,
+    "[#C][=C]": 496,
+    "[#C][=N]": 569,
+    "[#C][=O]": 515,
+    "[#C][C]": 407,
+    "[#C][N]": 465,
+    "[#C][O]": 538,
+    "[#C][S]": 534,
+    "[#N+1]": 123,
+    "[#N]": 19,
+    "[#N][=C]": 435,
+    "[#N][C]": 475,
+    "[#O+1]": 589,
+    "[#S]": 262,
+    "[-/Ring1]": 92,
+    "[-/Ring2]": 90,
+    "[-/Ring3]": 590,
+    "[-\\Ring1]": 167,
+    "[-\\Ring2]": 120,
+    "[-\\Ring3]": 591,
+    "[/123I]": 245,
+    "[/125I]": 117,
+    "[/131I]": 210,
+    "[/13CH1]": 272,
+    "[/13C]": 197,
+    "[/14CH1]": 324,
+    "[/14C]": 281,
+    "[/As]": 259,
+    "[/B]": 161,
+    "[/Br]": 62,
+    "[/C-1]": 229,
+    "[/C@@H1]": 45,
+    "[/C@@]": 105,
+    "[/C@H1]": 87,
+    "[/C@]": 76,
+    "[/C]": 20,
+    "[/C][=Branch1]": 428,
+    "[/C][=Branch2]": 450,
+    "[/C][=C]": 364,
+    "[/C][=N]": 454,
+    "[/C][Branch1]": 540,
+    "[/C][C]": 564,
+    "[/Cl]": 101,
+    "[/F]": 160,
+    "[/I]": 94,
+    "[/N+1]": 68,
+    "[/N-1]": 142,
+    "[/NH1]": 114,
+    "[/N]": 48,
+    "[/N][=C]": 418,
+    "[/N][=N]": 580,
+    "[/N][C]": 499,
+    "[/O-1]": 128,
+    "[/O]": 53,
+    "[/P+1]": 181,
+    "[/P@@]": 231,
+    "[/P]": 102,
+    "[/S+1]": 67,
+    "[/S-1]": 258,
+    "[/SH0]": 303,
+    "[/S]": 71,
+    "[/Se]": 149,
+    "[/Si]": 125,
+    "[10B]": 283,
+    "[11C@@H1]": 304,
+    "[11CH2]": 249,
+    "[11CH3]": 74,
+    "[11C]": 137,
+    "[123I-1]": 312,
+    "[123I]": 69,
+    "[123Te]": 186,
+    "[124I-1]": 325,
+    "[124I]": 133,
+    "[125I-1]": 171,
+    "[125I]": 70,
+    "[127I]": 301,
+    "[127Xe]": 311,
+    "[129Xe]": 292,
+    "[131Cs]": 275,
+    "[131I-1]": 329,
+    "[131I]": 86,
+    "[133Xe]": 322,
+    "[13CH1]": 198,
+    "[13CH3]": 277,
+    "[13C]": 166,
+    "[14C@@H1]": 318,
+    "[14C@@]": 302,
+    "[14C@H1]": 306,
+    "[14CH1]": 185,
+    "[14CH2]": 175,
+    "[14CH3]": 177,
+    "[14C]": 143,
+    "[15NH1]": 200,
+    "[15N]": 134,
+    "[17F]": 218,
+    "[18F-1]": 294,
+    "[18F]": 73,
+    "[18OH1]": 252,
+    "[19F]": 144,
+    "[211At]": 286,
+    "[223Ra]": 193,
+    "[22Na+1]": 191,
+    "[2H]": 88,
+    "[32P]": 307,
+    "[35S]": 182,
+    "[3H]": 103,
+    "[42K+1]": 188,
+    "[45Ca+2]": 178,
+    "[47Ca+2]": 201,
+    "[68Ga+3]": 273,
+    "[73Se]": 242,
+    "[75Se]": 327,
+    "[76Br]": 152,
+    "[81Kr]": 326,
+    "[82Rb+1]": 314,
+    "[82Rb]": 280,
+    "[85Sr+2]": 180,
+    "[89Sr+2]": 319,
+    "[=11C]": 132,
+    "[=13CH1]": 196,
+    "[=13C]": 199,
+    "[=14CH1]": 156,
+    "[=14C]": 194,
+    "[=18O]": 317,
+    "[=32PH1]": 257,
+    "[=32P]": 244,
+    "[=Al]": 295,
+    "[=As]": 145,
+    "[=B-1]": 159,
+    "[=B]": 176,
+    "[=Branch1]": 13,
+    "[=Branch1][=C]": 451,
+    "[=Branch1][=N]": 563,
+    "[=Branch1][=O]": 546,
+    "[=Branch1][C]": 335,
+    "[=Branch1][N+1]": 448,
+    "[=Branch1][N]": 508,
+    "[=Branch1][O]": 518,
+    "[=Branch2]": 14,
+    "[=Branch2][=C]": 472,
+    "[=Branch2][=N]": 584,
+    "[=Branch2][=O]": 549,
+    "[=Branch2][=Ring1]": 556,
+    "[=Branch2][C]": 383,
+    "[=Branch2][N]": 424,
+    "[=Branch2][O]": 507,
+    "[=Branch2][S]": 509,
+    "[=Branch3]": 592,
+    "[=C+1]": 267,
+    "[=CH0]": 593,
+    "[=C]": 7,
+    "[=C][=C]": 492,
+    "[=C][=O]": 516,
+    "[=C][Branch1]": 467,
+    "[=C][C]": 388,
+    "[=C][N]": 437,
+    "[=C][O]": 495,
+    "[=C][Ring1]": 482,
+    "[=C][S]": 456,
+    "[=Mg]": 310,
+    "[=N+1]": 32,
+    "[=N-1]": 63,
+    "[=NH1+1]": 232,
+    "[=NH2+1]": 243,
+    "[=N]": 29,
+    "[=N][=C]": 473,
+    "[=N][=N]": 579,
+    "[=N][=O]": 517,
+    "[=N][C]": 362,
+    "[=N][NH1]": 443,
+    "[=N][N]": 382,
+    "[=N][O]": 427,
+    "[=N][Ring1]": 513,
+    "[=N][S]": 501,
+    "[=O+1]": 116,
+    "[=OH0]": 296,
+    "[=OH1+1]": 594,
+    "[=O]": 18,
+    "[=O][/C]": 572,
+    "[=O][=C]": 411,
+    "[=O][=N]": 583,
+    "[=O][=O]": 425,
+    "[=O][C@@H1]": 414,
+    "[=O][C@@]": 544,
+    "[=O][C@H1]": 431,
+    "[=O][C@]": 514,
+    "[=O][C]": 361,
+    "[=O][NH1]": 434,
+    "[=O][N]": 345,
+    "[=O][O-1]": 391,
+    "[=O][O]": 370,
+    "[=O][S]": 571,
+    "[=P@@]": 107,
+    "[=P@]": 129,
+    "[=PH1]": 147,
+    "[=PH2]": 204,
+    "[=P]": 57,
+    "[=Ring1]": 34,
+    "[=Ring1][#Branch1]": 570,
+    "[=Ring1][#Branch2]": 576,
+    "[=Ring1][#C]": 562,
+    "[=Ring1][=Branch1]": 539,
+    "[=Ring1][=Branch2]": 481,
+    "[=Ring1][=C]": 582,
+    "[=Ring1][=N]": 545,
+    "[=Ring1][Branch1]": 442,
+    "[=Ring1][Branch2]": 568,
+    "[=Ring1][N]": 557,
+    "[=Ring1][O]": 526,
+    "[=Ring1][P]": 581,
+    "[=Ring1][S]": 565,
+    "[=Ring2]": 41,
+    "[=Ring2][Ring1]": 506,
+    "[=S+1]": 98,
+    "[=S-1]": 278,
+    "[=S@@]": 50,
+    "[=S@]": 164,
+    "[=SH0]": 269,
+    "[=SH1]": 150,
+    "[=S]": 26,
+    "[=S][N]": 455,
+    "[=Se+1]": 174,
+    "[=Se]": 162,
+    "[=Si]": 173,
+    "[=Te+1]": 179,
+    "[=Te]": 253,
+    "[Ag+1]": 190,
+    "[Ag-4]": 328,
+    "[Ag]": 146,
+    "[Al+3]": 223,
+    "[Al]": 225,
+    "[Ar]": 266,
+    "[As+1]": 189,
+    "[As-1]": 298,
+    "[AsH1]": 282,
+    "[AsH3]": 316,
+    "[As]": 84,
+    "[At]": 290,
+    "[Au-1]": 595,
+    "[B-1]": 100,
+    "[B@-1]": 206,
+    "[B@@-1]": 208,
+    "[BH0]": 165,
+    "[BH1-1]": 251,
+    "[BH2-1]": 202,
+    "[BH3-1]": 131,
+    "[B]": 65,
+    "[Ba+2]": 157,
+    "[Ba]": 291,
+    "[Be+2]": 246,
+    "[Bi+3]": 169,
+    "[Bi]": 155,
+    "[Br+2]": 254,
+    "[Br-1]": 52,
+    "[BrH0]": 596,
+    "[Br]": 21,
+    "[Br][=C]": 432,
+    "[Br][C]": 558,
+    "[Branch1]": 11,
+    "[Branch1][#Branch1]": 410,
+    "[Branch1][#Branch2]": 392,
+    "[Branch1][#C]": 399,
+    "[Branch1][=Branch1]": 401,
+    "[Branch1][=Branch2]": 358,
+    "[Branch1][=C]": 379,
+    "[Branch1][=N]": 405,
+    "[Branch1][=O]": 532,
+    "[Branch1][Branch1]": 476,
+    "[Branch1][Branch2]": 394,
+    "[Branch1][C]": 336,
+    "[Branch1][N]": 378,
+    "[Branch1][O]": 398,
+    "[Branch1][P]": 403,
+    "[Branch1][Ring1]": 396,
+    "[Branch1][Ring2]": 429,
+    "[Branch1][S]": 402,
+    "[Branch2]": 8,
+    "[Branch2][#Branch1]": 551,
+    "[Branch2][=Branch1]": 446,
+    "[Branch2][=C]": 464,
+    "[Branch2][=N]": 560,
+    "[Branch2][=O]": 522,
+    "[Branch2][Branch1]": 384,
+    "[Branch2][C]": 409,
+    "[Branch2][N]": 440,
+    "[Branch2][O]": 478,
+    "[Branch2][Ring1]": 340,
+    "[Branch2][Ring2]": 350,
+    "[Branch2][S]": 535,
+    "[Branch3]": 597,
+    "[C+1]": 106,
+    "[C-1]": 119,
+    "[C@@H1]": 27,
+    "[C@@H1][Branch1]": 371,
+    "[C@@H1][C@@H1]": 547,
+    "[C@@H1][C@H1]": 586,
+    "[C@@H1][C]": 447,
+    "[C@@H1][O]": 493,
+    "[C@@H1][Ring1]": 524,
+    "[C@@]": 28,
+    "[C@@][Branch1]": 413,
+    "[C@H1]": 33,
+    "[C@H1][Branch1]": 366,
+    "[C@H1][Branch2]": 573,
+    "[C@H1][C@@H1]": 505,
+    "[C@H1][C@H1]": 575,
+    "[C@H1][C]": 445,
+    "[C@H1][O]": 530,
+    "[C@H1][Ring1]": 470,
+    "[C@]": 35,
+    "[C@][Branch1]": 412,
+    "[CH0]": 227,
+    "[CH1-1]": 138,
+    "[CH1]": 323,
+    "[CH2-1]": 598,
+    "[CH2]": 224,
+    "[CH3]": 226,
+    "[C]": 5,
+    "[C][#C]": 500,
+    "[C][#N]": 536,
+    "[C][=Branch1]": 377,
+    "[C][=Branch2]": 488,
+    "[C][=C]": 333,
+    "[C][=N+1]": 543,
+    "[C][=N]": 346,
+    "[C][=O]": 511,
+    "[C][=Ring1]": 395,
+    "[C][=Ring2]": 471,
+    "[C][Branch1]": 342,
+    "[C][Branch2]": 423,
+    "[C][C@@H1]": 386,
+    "[C][C@@]": 480,
+    "[C][C@H1]": 385,
+    "[C][C@]": 462,
+    "[C][C]": 334,
+    "[C][Cl]": 574,
+    "[C][F]": 578,
+    "[C][NH1]": 449,
+    "[C][N]": 352,
+    "[C][O]": 348,
+    "[C][Ring1]": 351,
+    "[C][Ring2]": 588,
+    "[C][S]": 397,
+    "[Ca+2]": 122,
+    "[CaH2]": 288,
+    "[Ca]": 234,
+    "[Cl+1]": 139,
+    "[Cl+2]": 240,
+    "[Cl+3]": 61,
+    "[Cl-1]": 39,
+    "[ClH0]": 126,
+    "[ClH1+1]": 599,
+    "[ClH2+1]": 600,
+    "[Cl]": 24,
+    "[Cl][=C]": 374,
+    "[Cl][=N]": 494,
+    "[Cl][C]": 416,
+    "[Cl][Cl]": 484,
+    "[Cl][N]": 512,
+    "[Cl][O]": 587,
+    "[Cl][S]": 555,
+    "[Co-2]": 601,
+    "[Co-4]": 602,
+    "[Cr]": 603,
+    "[Cs+1]": 85,
+    "[Cs]": 284,
+    "[Cu-2]": 604,
+    "[Cu-3]": 605,
+    "[Cu-5]": 606,
+    "[Cu]": 607,
+    "[F-1]": 111,
+    "[F]": 16,
+    "[F][=C]": 376,
+    "[F][=N]": 537,
+    "[F][C]": 404,
+    "[F][F]": 360,
+    "[F][N]": 585,
+    "[F][O]": 577,
+    "[Fe+2]": 608,
+    "[Fe+3]": 609,
+    "[Fe-2]": 610,
+    "[Fe-3]": 611,
+    "[Gd+3]": 612,
+    "[Ge]": 613,
+    "[H+1]": 151,
+    "[H-1]": 285,
+    "[H]": 614,
+    "[He]": 115,
+    "[Hg]": 615,
+    "[I+1]": 153,
+    "[I+2]": 308,
+    "[I+3]": 309,
+    "[I-1]": 66,
+    "[IH0]": 183,
+    "[I]": 42,
+    "[Ir-3]": 616,
+    "[K+1]": 83,
+    "[KH1]": 228,
+    "[K]": 313,
+    "[Kr]": 315,
+    "[Li+1]": 110,
+    "[LiH1]": 233,
+    "[Li]": 276,
+    "[Mg+1]": 289,
+    "[Mg+2]": 75,
+    "[MgH2]": 236,
+    "[Mg]": 235,
+    "[Mn+1]": 617,
+    "[Mn+2]": 618,
+    "[Mo]": 619,
+    "[N+1]": 30,
+    "[N+1][=C]": 510,
+    "[N+1][Branch1]": 497,
+    "[N-1]": 37,
+    "[N@+1]": 97,
+    "[N@@+1]": 109,
+    "[N@@]": 140,
+    "[N@H1+1]": 297,
+    "[N@]": 91,
+    "[NH0]": 238,
+    "[NH1+1]": 216,
+    "[NH1-1]": 124,
+    "[NH1]": 38,
+    "[NH1][C]": 406,
+    "[NH1][Ring1]": 566,
+    "[NH2+1]": 215,
+    "[NH2-1]": 620,
+    "[NH2]": 621,
+    "[NH3+1]": 293,
+    "[NH4+1]": 260,
+    "[N]": 10,
+    "[N][#C]": 422,
+    "[N][=C]": 343,
+    "[N][=N]": 389,
+    "[N][=O]": 417,
+    "[N][=Ring1]": 420,
+    "[N][=Ring2]": 457,
+    "[N][Branch1]": 347,
+    "[N][Branch2]": 498,
+    "[N][C@@H1]": 486,
+    "[N][C@H1]": 483,
+    "[N][C]": 339,
+    "[N][N]": 430,
+    "[N][O]": 468,
+    "[N][Ring1]": 415,
+    "[N][S]": 408,
+    "[Na+1]": 40,
+    "[NaH1]": 221,
+    "[Na]": 268,
+    "[Ni-2]": 622,
+    "[Ni-4]": 623,
+    "[O+1]": 136,
+    "[O-1]": 31,
+    "[O-2]": 270,
+    "[OH0]": 96,
+    "[OH1+1]": 220,
+    "[OH1-1]": 154,
+    "[OH1]": 212,
+    "[OH2+1]": 624,
+    "[OH3+1]": 274,
+    "[O]": 6,
+    "[O][=C]": 349,
+    "[O][=N+1]": 503,
+    "[O][=N]": 474,
+    "[O][=O]": 521,
+    "[O][=S]": 460,
+    "[O][Branch1]": 559,
+    "[O][C@@H1]": 390,
+    "[O][C@@]": 553,
+    "[O][C@H1]": 387,
+    "[O][C@]": 519,
+    "[O][C]": 341,
+    "[O][N]": 441,
+    "[O][O]": 426,
+    "[O][P]": 489,
+    "[O][Ring1]": 458,
+    "[O][S]": 487,
+    "[P+1]": 112,
+    "[P-1]": 135,
+    "[P@+1]": 163,
+    "[P@@+1]": 237,
+    "[P@@]": 78,
+    "[P@]": 72,
+    "[PH1]": 127,
+    "[PH2+1]": 305,
+    "[P]": 25,
+    "[P][=C]": 504,
+    "[P][=O]": 552,
+    "[P][C]": 419,
+    "[P][N]": 477,
+    "[P][O]": 548,
+    "[P][S]": 533,
+    "[Pd-2]": 625,
+    "[Pt-2]": 626,
+    "[Pt]": 627,
+    "[Ra]": 187,
+    "[Rb+1]": 207,
+    "[Rb]": 321,
+    "[Rh+2]": 628,
+    "[Ring1]": 9,
+    "[Ring1][#Branch1]": 344,
+    "[Ring1][#Branch2]": 353,
+    "[Ring1][#C]": 369,
+    "[Ring1][=Branch1]": 338,
+    "[Ring1][=Branch2]": 356,
+    "[Ring1][=C]": 363,
+    "[Ring1][=N]": 367,
+    "[Ring1][=O]": 550,
+    "[Ring1][Branch1]": 355,
+    "[Ring1][Branch2]": 357,
+    "[Ring1][C]": 373,
+    "[Ring1][N]": 368,
+    "[Ring1][O]": 354,
+    "[Ring1][P]": 372,
+    "[Ring1][Ring1]": 421,
+    "[Ring1][Ring2]": 439,
+    "[Ring1][S]": 365,
+    "[Ring2]": 12,
+    "[Ring2][=C]": 381,
+    "[Ring2][=N]": 459,
+    "[Ring2][=O]": 525,
+    "[Ring2][Branch1]": 520,
+    "[Ring2][C]": 393,
+    "[Ring2][N]": 433,
+    "[Ring2][O]": 466,
+    "[Ring2][Ring1]": 337,
+    "[Ring2][Ring2]": 359,
+    "[Ring2][S]": 502,
+    "[Ring3]": 629,
+    "[S+1]": 44,
+    "[S-1]": 60,
+    "[S-2]": 172,
+    "[S@+1]": 56,
+    "[S@@+1]": 95,
+    "[S@@]": 51,
+    "[S@]": 184,
+    "[SH0]": 170,
+    "[SH1+1]": 263,
+    "[SH1-1]": 205,
+    "[SH1]": 264,
+    "[SH2]": 299,
+    "[S]": 17,
+    "[S][=C]": 490,
+    "[S][=N]": 523,
+    "[S][=O]": 554,
+    "[S][Branch1]": 444,
+    "[S][C]": 375,
+    "[S][N]": 452,
+    "[S][O]": 527,
+    "[S][Ring1]": 485,
+    "[S][S]": 463,
+    "[Sb]": 211,
+    "[Se+1]": 230,
+    "[SeH1]": 130,
+    "[SeH2]": 330,
+    "[Se]": 36,
+    "[Si-1]": 331,
+    "[Si@]": 287,
+    "[SiH1-1]": 192,
+    "[SiH1]": 104,
+    "[SiH2]": 300,
+    "[SiH3-1]": 219,
+    "[SiH4]": 271,
+    "[Si]": 49,
+    "[Sn]": 630,
+    "[Sr+2]": 158,
+    "[Sr]": 332,
+    "[TeH1]": 279,
+    "[TeH2]": 239,
+    "[Te]": 81,
+    "[Xe]": 320,
+    "[Yb]": 209,
+    "[Zn+1]": 213,
+    "[Zn+2]": 148,
+    "[Zn-2]": 261,
+    "[Zn]": 82,
+    "[\\11CH3]": 203,
+    "[\\123I]": 214,
+    "[\\125I]": 79,
+    "[\\3H]": 250,
+    "[\\B]": 217,
+    "[\\Br]": 89,
+    "[\\C-1]": 265,
+    "[\\C@@H1]": 46,
+    "[\\C@@]": 80,
+    "[\\C@H1]": 77,
+    "[\\C@]": 99,
+    "[\\CH1-1]": 241,
+    "[\\C]": 43,
+    "[\\C][=C]": 436,
+    "[\\C][C]": 567,
+    "[\\Cl]": 58,
+    "[\\F]": 118,
+    "[\\I]": 59,
+    "[\\N+1]": 113,
+    "[\\N-1]": 222,
+    "[\\NH1]": 108,
+    "[\\N]": 55,
+    "[\\N][C]": 541,
+    "[\\O-1]": 64,
+    "[\\O]": 47,
+    "[\\P]": 93,
+    "[\\S+1]": 121,
+    "[\\S-1]": 255,
+    "[\\S]": 54,
+    "[\\Se]": 168,
+    "[\\Si]": 195
+}

ape_tokenizer/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>",
+  "mask_token": "<mask>"
+}

ape_tokenizer/tokenization_ape.py

@@ -0,0 +1,709 @@
+"""Hugging Face-compatible tokenizer for APE molecular vocabularies.
+
+This file is intentionally self-contained so it can be copied into a model repo
+and loaded by ``AutoTokenizer.from_pretrained(..., trust_remote_code=True)``.
+"""
+
+import json
+import os
+import re
+from collections.abc import Mapping
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Literal
+
+from transformers import PreTrainedTokenizer
+
+
+Representation = Literal["SELFIES", "SMILES"]
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "selfies_vocab_file": "selfies_vocab.json",
+    "smiles_vocab_file": "smiles_vocab.json",
+}
+SELFIES_RE = re.compile(r"\[[^\]]+\]")
+SMILES_RE = re.compile(
+    r"(\[[^\]]+\]|Br?|Cl?|Si?|Se?|Li?|Na?|Mg?|Al?|Ca?|Fe?|Zn?|"
+    r"N|O|S|P|F|I|K|B|C|H|"
+    r"b|c|n|o|s|p|"
+    r"\%\d{2}|\d|"
+    r"\(|\)|\.|=|#|-|\+|\\|/|:|~|@|\?|\*|\$)"
+)
+
+
+def _base_piece_count(token: str, representation: str) -> int:
+    """Count primitive molecular pieces in a vocab token."""
+    pieces = pre_tokenize_molecule(token, representation)
+    return max(1, len(pieces))
+
+
+def _max_vocab_piece_span(vocab: dict[str, int], representation: str) -> int:
+    """Maximum number of primitive pieces covered by any non-special vocab token."""
+    max_span = 1
+    for token in vocab:
+        if token.startswith("<") and token.endswith(">"):
+            continue
+        max_span = max(max_span, _base_piece_count(token, representation))
+    return max_span
+
+
+def _coerce_vocab(vocab: Mapping[str, Any]) -> dict[str, int]:
+    if not isinstance(vocab, Mapping):
+        raise ValueError("Vocabulary must be a JSON object mapping token strings to integer IDs.")
+    out = {str(token): int(idx) for token, idx in vocab.items()}
+    if len(set(out.values())) != len(out):
+        raise ValueError("Vocabulary token IDs must be unique.")
+    return out
+
+
+def _token_text(token: Any) -> str:
+    return str(getattr(token, "content", token))
+
+
+def _normalize_representation(representation: str) -> Representation:
+    normalized = representation.upper()
+    if normalized not in {"SELFIES", "SMILES"}:
+        raise ValueError(f"representation must be 'SELFIES' or 'SMILES', got {representation!r}")
+    return normalized  # type: ignore[return-value]
+
+
+def _select_vocab_file(
+    *,
+    representation: Representation,
+    vocab_file: str | os.PathLike[str] | None,
+    selfies_vocab_file: str | os.PathLike[str] | None,
+    smiles_vocab_file: str | os.PathLike[str] | None,
+) -> str | os.PathLike[str] | None:
+    if representation == "SELFIES" and selfies_vocab_file is not None:
+        return selfies_vocab_file
+    if representation == "SMILES" and smiles_vocab_file is not None:
+        return smiles_vocab_file
+    return vocab_file
+
+
+def pre_tokenize_molecule(molecule: str, representation: str) -> list[str]:
+    active_representation = _normalize_representation(representation)
+    if active_representation == "SELFIES":
+        return SELFIES_RE.findall(molecule)
+
+    tokens: list[str] = []
+    cursor = 0
+    for match in SMILES_RE.finditer(molecule):
+        if match.start() > cursor:
+            tokens.extend(molecule[cursor : match.start()])
+        tokens.append(match.group(0))
+        cursor = match.end()
+    if cursor < len(molecule):
+        tokens.extend(molecule[cursor:])
+    return [token for token in tokens if token and not token.isspace()]
+
+
+def ape_tokenize(
+    text: str,
+    vocab: dict[str, int],
+    representation: str,
+    unk_token: str = "<unk>",
+    max_piece_span: int | None = None,
+) -> list[str]:
+    pieces = pre_tokenize_molecule(text, representation)
+    if not pieces:
+        return [unk_token]
+
+    if max_piece_span is None:
+        max_piece_span = _max_vocab_piece_span(vocab, representation)
+
+    n = len(pieces)
+    tokens: list[str] = []
+    append_token = tokens.append
+    vocab_contains = vocab.__contains__
+    join_pieces = "".join
+    i = 0
+
+    while i < n:
+        upper = min(n, i + max_piece_span)
+
+        for j in range(upper, i, -1):
+            candidate = join_pieces(pieces[i:j])
+            if vocab_contains(candidate):
+                append_token(candidate)
+                i = j
+                break
+        else:
+            append_token(unk_token)
+            i += 1
+
+    return tokens
+
+
+class APEPreTrainedTokenizer(PreTrainedTokenizer):
+    """Hugging Face tokenizer backend for APE molecular tokenization. (Not fast)"""
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str | os.PathLike[str] | None = None,
+        selfies_vocab_file: str | os.PathLike[str] | None = None,
+        smiles_vocab_file: str | os.PathLike[str] | None = None,
+        vocab: dict[str, Any] | None = None,
+        representation: str = "SELFIES",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+        unk_token: str = "<unk>",
+        pad_token: str = "<pad>",
+        mask_token: str = "<mask>",
+        model_max_length: int = 256,
+        **kwargs,
+    ) -> None:
+        self.representation = _normalize_representation(representation)
+        active_vocab_file = _select_vocab_file(
+            representation=self.representation,
+            vocab_file=vocab_file,
+            selfies_vocab_file=selfies_vocab_file,
+            smiles_vocab_file=smiles_vocab_file,
+        )
+
+        if vocab is None:
+            if active_vocab_file is None:
+                vocab = {
+                    bos_token: 0,
+                    pad_token: 1,
+                    eos_token: 2,
+                    unk_token: 3,
+                    mask_token: 4,
+                }
+            else:
+                with open(active_vocab_file, encoding="utf-8") as f:
+                    vocab = json.load(f)
+
+        if vocab is None:
+            raise ValueError("Loaded vocabulary is None.")
+
+        self.vocab_file = str(active_vocab_file) if active_vocab_file is not None else None
+        self.selfies_vocab_file = (
+            str(selfies_vocab_file) if selfies_vocab_file is not None else None
+        )
+        self.smiles_vocab_file = str(smiles_vocab_file) if smiles_vocab_file is not None else None
+        self.vocab = _coerce_vocab(vocab)
+        self._require_special_tokens(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+        )
+        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
+        self.vocabulary_frequency: dict[str, int] = {}
+        self.pair_counts: dict[tuple[str, str], int] = {}
+        self._max_piece_span = _max_vocab_piece_span(self.vocab, self.representation)
+
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            model_max_length=model_max_length,
+            representation=self.representation,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    @property
+    def vocabulary(self) -> dict[str, int]:
+        """Legacy alias for callers that previously used APETokenizer."""
+        return self.vocab
+
+    @vocabulary.setter
+    def vocabulary(self, value: dict[str, int]) -> None:
+        self.vocab = _coerce_vocab(value)
+        self.update_reverse_vocabulary()
+        self._refresh_tokenization_cache()
+
+    @property
+    def special_tokens(self) -> dict[str, int]:
+        bos_token = str(self.bos_token)
+        pad_token = str(self.pad_token)
+        eos_token = str(self.eos_token)
+        unk_token = str(self.unk_token)
+        mask_token = str(self.mask_token)
+        return {
+            bos_token: self._convert_token_to_id(bos_token),
+            pad_token: self._convert_token_to_id(pad_token),
+            eos_token: self._convert_token_to_id(eos_token),
+            unk_token: self._convert_token_to_id(unk_token),
+            mask_token: self._convert_token_to_id(mask_token),
+        }
+
+    @special_tokens.setter
+    def special_tokens(self, value: dict[str, int]) -> None:
+        for token, token_id in value.items():
+            self.vocab.setdefault(str(token), int(token_id))
+        self.vocab = _coerce_vocab(self.vocab)
+        self.update_reverse_vocabulary()
+        self._refresh_tokenization_cache()
+
+    def get_vocab(self) -> dict[str, int]:
+        return dict(self.vocab)
+
+    def update_reverse_vocabulary(self) -> None:
+        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
+
+    def _refresh_tokenization_cache(self) -> None:
+        self._max_piece_span = _max_vocab_piece_span(self.vocab, self.representation)
+
+    def _require_special_tokens(
+        self,
+        *,
+        bos_token: str,
+        eos_token: str,
+        unk_token: str,
+        pad_token: str,
+        mask_token: str,
+    ) -> None:
+        missing = [
+            token_text
+            for token in [bos_token, eos_token, unk_token, pad_token, mask_token]
+            if (token_text := _token_text(token)) not in self.vocab
+        ]
+        if missing:
+            raise ValueError(f"Vocabulary is missing required special tokens: {missing}")
+
+    def pre_tokenize(self, molecule: str, representation: str | None = None) -> list[str]:
+        return pre_tokenize_molecule(molecule, representation or self.representation)
+
+    def _tokenize(self, text: str, **kwargs) -> list[str]:
+
+        return ape_tokenize(
+            text,
+            vocab=self.vocab,
+            representation=self.representation,
+            unk_token=str(self.unk_token),
+            max_piece_span=self._max_piece_span,
+        )
+
+    def encode_molecule(
+        self,
+        text: str,
+        add_special_tokens: bool = True,
+        max_length: int | None = None,
+        truncation: bool = True,
+    ) -> list[int]:
+        """Fast molecular encode path avoiding generic Hugging Face tokenizer overhead."""
+
+        tokens = self._tokenize(text)
+
+        ids = [self._convert_token_to_id(token) for token in tokens]
+
+        if add_special_tokens:
+            ids = self.build_inputs_with_special_tokens(ids)
+
+        if max_length is not None and truncation:
+            ids = ids[:max_length]
+
+        return ids
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.vocab.get(token, self.vocab[str(self.unk_token)])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.ids_to_tokens.get(int(index), str(self.unk_token))
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return "".join(tokens)
+
+    def _required_special_token_id(
+        self,
+        token_value: int | list[int] | str | list[str] | None,
+        token_name: str,
+    ) -> int:
+        if token_value is None:
+            raise ValueError(f"{token_name} must be set.")
+        if isinstance(token_value, int):
+            return token_value
+        if isinstance(token_value, str):
+            return self._convert_token_to_id(token_value)
+        if len(token_value) == 1:
+            only_value = token_value[0]
+            if isinstance(only_value, int):
+                return only_value
+            if isinstance(only_value, str):
+                return self._convert_token_to_id(only_value)
+        raise ValueError(f"{token_name} must resolve to a single token id.")
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        bos_id = self._required_special_token_id(self.bos_token, "bos_token")
+        eos_id = self._required_special_token_id(self.eos_token, "eos_token")
+        if token_ids_1 is None:
+            return [bos_id, *token_ids_0, eos_id]
+        return [bos_id, *token_ids_0, eos_id, *token_ids_1, eos_id]
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1))
+
+    def pad(
+        self,
+        encoded_inputs: Any,
+        padding: Any = True,
+        max_length: int | None = None,
+        pad_to_multiple_of: int | None = None,
+        padding_side: str | None = None,
+        return_attention_mask: bool | None = None,
+        return_tensors: Any = None,
+        verbose: bool = True,
+    ):
+        padding_enabled = padding not in (False, "do_not_pad")
+        if (
+            padding_enabled
+            and isinstance(encoded_inputs, list)
+            and any("labels" in item for item in encoded_inputs)
+        ):
+            target_length = max(
+                len(item.get("input_ids", item.get("labels", []))) for item in encoded_inputs
+            )
+            if padding == "max_length" and max_length is not None:
+                target_length = max_length
+
+            if pad_to_multiple_of and target_length % pad_to_multiple_of:
+                target_length = ((target_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
+
+            padded_inputs = []
+            for item in encoded_inputs:
+                item = dict(item)
+                labels = list(item.get("labels", []))
+                pad_len = max(0, target_length - len(labels))
+                if pad_len:
+                    label_padding = [-100] * pad_len
+                    if self.padding_side == "left":
+                        labels = label_padding + labels
+                    else:
+                        labels = labels + label_padding
+                    item["labels"] = labels
+                padded_inputs.append(item)
+            encoded_inputs = padded_inputs
+
+        return super().pad(
+            encoded_inputs,
+            padding=padding,
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            padding_side=padding_side,
+            return_attention_mask=return_attention_mask,
+            return_tensors=return_tensors,
+            verbose=verbose,
+        )
+
+    def save_vocabulary(
+        self,
+        save_directory: str,
+        filename_prefix: str | None = None,
+    ) -> tuple[str, ...]:
+        if not os.path.isdir(save_directory):
+            raise ValueError(f"Vocabulary path ({save_directory}) should be a directory.")
+
+        vocab_file = Path(save_directory) / (
+            f"{filename_prefix}-vocab.json" if filename_prefix else "vocab.json"
+        )
+        with vocab_file.open("w", encoding="utf-8") as f:
+            json.dump(self.vocab, f, ensure_ascii=False, indent=4)
+        return (str(vocab_file),)
+
+    def add_tokens_to_vocabulary(self, tokens: list[str]) -> int:
+        """Add tokens to the tokenizer vocabulary if they are not already present.
+
+        This is intended for forcing coverage of rare valid molecular primitive
+        symbols, especially SELFIES bracket tokens, after APE merge training.
+        """
+
+        if not tokens:
+            return 0
+
+        next_id = max(self.vocab.values(), default=-1) + 1
+        added = 0
+
+        for token in tokens:
+            token = str(token).strip()
+            if not token:
+                continue
+            if token in self.vocab:
+                continue
+
+            self.vocab[token] = next_id
+            next_id += 1
+            added += 1
+
+        if added:
+            self.update_reverse_vocabulary()
+            self._refresh_tokenization_cache()
+
+        return added
+
+    def save_pretrained(self, save_directory: str | os.PathLike[str], *args, **kwargs):
+        saved_files = super().save_pretrained(save_directory, *args, **kwargs)
+        save_path = Path(save_directory)
+
+        special_tokens_map = {
+            "bos_token": str(self.bos_token),
+            "eos_token": str(self.eos_token),
+            "unk_token": str(self.unk_token),
+            "pad_token": str(self.pad_token),
+            "mask_token": str(self.mask_token),
+        }
+        with (save_path / "special_tokens_map.json").open("w", encoding="utf-8") as f:
+            json.dump(special_tokens_map, f, ensure_ascii=False, indent=2)
+
+        tokenizer_config_path = save_path / "tokenizer_config.json"
+        if tokenizer_config_path.exists():
+            with tokenizer_config_path.open(encoding="utf-8") as f:
+                tokenizer_config = json.load(f)
+        else:
+            tokenizer_config = {}
+        tokenizer_config.pop("tokenizer_class", None)
+        tokenizer_config.update(
+            {
+                "representation": self.representation,
+                "model_max_length": self.model_max_length,
+                "auto_map": {
+                    "AutoTokenizer": [
+                        "tokenization_ape.APEPreTrainedTokenizer",
+                        None,
+                    ],
+                },
+            }
+        )
+        with tokenizer_config_path.open("w", encoding="utf-8") as f:
+            json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)
+
+        return saved_files
+
+    def save_vocabulary_file(self, file_path: str | os.PathLike[str]) -> None:
+        path = Path(file_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        freq_path = path.with_name(f"{path.stem}_freq.json")
+
+        with path.open("w", encoding="utf-8") as f:
+            json.dump(self.vocab, f, ensure_ascii=False, indent=4)
+        with freq_path.open("w", encoding="utf-8") as f:
+            json.dump(self.vocabulary_frequency, f, ensure_ascii=False, indent=4)
+
+    def load_vocabulary_file(
+        self,
+        file_path: str | os.PathLike[str],
+        representation: str | None = None,
+    ) -> None:
+        if representation is not None:
+            self.representation = _normalize_representation(representation)
+        with open(file_path, encoding="utf-8") as f:
+            vocab = json.load(f)
+        self.vocab = _coerce_vocab(vocab)
+        self._require_special_tokens(
+            bos_token=str(self.bos_token),
+            eos_token=str(self.eos_token),
+            unk_token=str(self.unk_token),
+            pad_token=str(self.pad_token),
+            mask_token=str(self.mask_token),
+        )
+        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
+        self._refresh_tokenization_cache()
+
+    def train(
+        self,
+        corpus,
+        type: str = "selfies",
+        representation: str | None = None,
+        max_vocab_size: int = 5000,
+        min_freq_for_merge: int = 2000,
+        max_merge_pieces: int | None = 8,
+        save_checkpoint: bool = False,
+        checkpoint_path: str = "checkpoint",
+        checkpoint_interval: int = 500,
+    ) -> None:
+        import warnings
+
+        new_rep = _normalize_representation(representation or type)
+        if new_rep != self.representation:
+            warnings.warn(
+                f"train() representation={new_rep!r} differs from tokenizer "
+                f"representation={self.representation!r}. Overwriting.",
+                UserWarning,
+                stacklevel=2,
+            )
+        self.representation = new_rep
+
+        if not corpus:
+            raise ValueError("Cannot train APE tokenizer on an empty corpus.")
+
+        print(f"Pretokenizing {self.representation}...", flush=True)
+        tokenized_corpus = []
+        vocabulary_frequency: defaultdict[str, int] = defaultdict(int)
+        saw_tokens = False
+
+        for sentence in corpus:
+            tokens = self.pre_tokenize(str(sentence))
+            if not tokens:
+                continue
+            saw_tokens = True
+            for token in tokens:
+                vocabulary_frequency[token] += 1
+            if len(tokens) > 1:
+                tokenized_corpus.append(tokens)
+        print(
+            f"Pretokenization complete, found {len(vocabulary_frequency)} tokens",
+            flush=True,
+        )
+
+        if not saw_tokens:
+            raise ValueError("Cannot train APE tokenizer on an empty corpus.")
+
+        pre_tokens_counts = len(vocabulary_frequency)
+        merged_counter = len(vocabulary_frequency) + 1
+        if save_checkpoint and checkpoint_interval <= 0:
+            raise ValueError(
+                "checkpoint_interval must be positive when save_checkpoint is enabled."
+            )
+        checkpoint_increment = checkpoint_interval
+        batch = checkpoint_interval + pre_tokens_counts
+        piece_count_cache: dict[str, int] = {}
+
+        def merged_piece_count(token: str) -> int:
+            count = piece_count_cache.get(token)
+            if count is None:
+                count = _base_piece_count(token, self.representation)
+                piece_count_cache[token] = count
+            return count
+
+        def get_most_common_pair(tokenized):
+            pair_counts: defaultdict[tuple[str, str], int] = defaultdict(int)
+            for tokens in tokenized:
+                for i in range(len(tokens) - 1):
+                    pair = (tokens[i], tokens[i + 1])
+
+                    if max_merge_pieces is not None:
+                        merged_candidate = pair[0] + pair[1]
+                        if merged_piece_count(merged_candidate) > max_merge_pieces:
+                            continue
+
+                    pair_counts[pair] += 1
+
+            self.pair_counts = dict(pair_counts)
+            if not pair_counts:
+                return ("", ""), 0
+
+            most_common_pair = ("", "")
+            most_common_frequency = 0
+            for pair, count in pair_counts.items():
+                if count > most_common_frequency:
+                    most_common_pair = pair
+                    most_common_frequency = count
+            return most_common_pair, most_common_frequency
+
+        while True:
+            if save_checkpoint and len(vocabulary_frequency) >= batch:
+                self.vocabulary_frequency = dict(vocabulary_frequency)
+                self.vocab = {
+                    **{
+                        str(self.bos_token): 0,
+                        str(self.pad_token): 1,
+                        str(self.eos_token): 2,
+                        str(self.unk_token): 3,
+                        str(self.mask_token): 4,
+                    },
+                    **{
+                        word: idx
+                        for idx, word in enumerate(
+                            vocabulary_frequency.keys(),
+                            start=5,
+                        )
+                    },
+                }
+                self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
+                self._refresh_tokenization_cache()
+                checkpoint_dir = Path(checkpoint_path)
+                checkpoint_dir.mkdir(parents=True, exist_ok=True)
+                self.save_vocabulary_file(checkpoint_dir / f"checkpoint_{batch}.json")
+                self.save_pretrained(str(checkpoint_dir / f"checkpoint_{batch}"))
+                print(f"Checkpoint saved at {checkpoint_dir}/checkpoint_{batch}.json")
+                batch += checkpoint_increment
+
+            if len(vocabulary_frequency) >= max_vocab_size:
+                print("Max vocabulary achieved", flush=True)
+                break
+
+            if not tokenized_corpus:
+                print("No more mergeable pairs", flush=True)
+                break
+
+            most_common_pair, freq = get_most_common_pair(tokenized_corpus)
+            if freq < min_freq_for_merge:
+                print("Not enough frequency found", flush=True)
+                break
+
+            if not most_common_pair[0] or not most_common_pair[1]:
+                print("No valid merge pair found", flush=True)
+                break
+
+            left_token, right_token = most_common_pair
+            merged_word = left_token + right_token
+            if merged_word not in vocabulary_frequency:
+                print(
+                    f"New merge found: {merged_word} {merged_counter}/{max_vocab_size} "
+                    f"{round(merged_counter / max_vocab_size * 100, 2)}%",
+                    flush=True,
+                )
+                merged_counter += 1
+            vocabulary_frequency[merged_word] += freq
+
+            new_tokenized_corpus = []
+            for tokens in tokenized_corpus:
+                new_tokens = []
+                append_token = new_tokens.append
+                i = 0
+                token_count = len(tokens)
+                while i < token_count:
+                    if (
+                        i < token_count - 1
+                        and tokens[i] == left_token
+                        and tokens[i + 1] == right_token
+                    ):
+                        append_token(merged_word)
+                        i += 2
+                    else:
+                        append_token(tokens[i])
+                        i += 1
+
+                if len(new_tokens) > 1:
+                    new_tokenized_corpus.append(new_tokens)
+
+            tokenized_corpus = new_tokenized_corpus
+
+        self.vocabulary_frequency = dict(vocabulary_frequency)
+        self.vocab = {
+            str(self.bos_token): 0,
+            str(self.pad_token): 1,
+            str(self.eos_token): 2,
+            str(self.unk_token): 3,
+            str(self.mask_token): 4,
+            **{word: idx for idx, word in enumerate(vocabulary_frequency.keys(), start=5)},
+        }
+
+        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
+        self._refresh_tokenization_cache()
+
+    def train_from_iterator(self, iterator, *args, **kwargs) -> None:
+        raise NotImplementedError("train_from_iterator is not implemented for APE")
+
+
+APEPreTrainedTokenizer.register_for_auto_class("AutoTokenizer")

ape_tokenizer/tokenizer_config.json

@@ -0,0 +1,63 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "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": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_ape.APEPreTrainedTokenizer",
+      null
+    ]
+  },
+  "backend": "custom",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 256,
+  "pad_token": "<pad>",
+  "representation": "SELFIES",
+  "unk_token": "<unk>",
+  "use_fast": false
+}

ape_tokenizer/tokenizer_metadata.json

@@ -0,0 +1,28 @@
+{
+  "ape_source": "modernmolbert.local",
+  "created_at_utc": "2026-05-19T16:24:35.971419+00:00",
+  "creation_command": "python -m modernmolbert.train_ape_tokenizer",
+  "dataset_name": "data/pretrain/chembl36_selfies",
+  "extra_vocab_selfies_path": null,
+  "extra_vocab_symbols_added": 42,
+  "extra_vocab_symbols_path": "tokenizer/extra_symbols/benchmark_missing_selfies_symbols_min10.txt",
+  "extra_vocab_symbols_requested": 42,
+  "max_merge_pieces": 2,
+  "max_vocab_size": 2000,
+  "min_freq_for_merge": 3000,
+  "representation": "SELFIES",
+  "seed": 42,
+  "selfies_column": "selfies",
+  "shuffle_buffer_size": 100000,
+  "special_ids": {
+    "bos_token": 0,
+    "eos_token": 2,
+    "mask_token": 4,
+    "pad_token": 1,
+    "unk_token": 3
+  },
+  "tokenizer_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+  "tokenizer_sha256": "26ad0e90de9c0a469eb6e3b7aa985e2a2d7cc3f0dd7b2b229ccfce41f639e208",
+  "tokenizer_train_size": 2000000,
+  "vocab_size": 631
+}

ape_tokenizer/vocab.json

@@ -0,0 +1,633 @@
+{
+    "</s>": 2,
+    "<mask>": 4,
+    "<pad>": 1,
+    "<s>": 0,
+    "<unk>": 3,
+    "[#11C]": 256,
+    "[#14C]": 247,
+    "[#Branch1]": 22,
+    "[#Branch1][=C]": 479,
+    "[#Branch1][=O]": 542,
+    "[#Branch1][C]": 400,
+    "[#Branch1][N]": 453,
+    "[#Branch1][O]": 491,
+    "[#Branch1][S]": 528,
+    "[#Branch2]": 15,
+    "[#Branch2][=C]": 469,
+    "[#Branch2][=N]": 561,
+    "[#Branch2][=O]": 531,
+    "[#Branch2][C]": 380,
+    "[#Branch2][N]": 438,
+    "[#Branch2][O]": 461,
+    "[#Branch2][S]": 529,
+    "[#C-1]": 141,
+    "[#CH0]": 248,
+    "[#C]": 23,
+    "[#C][=C]": 496,
+    "[#C][=N]": 569,
+    "[#C][=O]": 515,
+    "[#C][C]": 407,
+    "[#C][N]": 465,
+    "[#C][O]": 538,
+    "[#C][S]": 534,
+    "[#N+1]": 123,
+    "[#N]": 19,
+    "[#N][=C]": 435,
+    "[#N][C]": 475,
+    "[#O+1]": 589,
+    "[#S]": 262,
+    "[-/Ring1]": 92,
+    "[-/Ring2]": 90,
+    "[-/Ring3]": 590,
+    "[-\\Ring1]": 167,
+    "[-\\Ring2]": 120,
+    "[-\\Ring3]": 591,
+    "[/123I]": 245,
+    "[/125I]": 117,
+    "[/131I]": 210,
+    "[/13CH1]": 272,
+    "[/13C]": 197,
+    "[/14CH1]": 324,
+    "[/14C]": 281,
+    "[/As]": 259,
+    "[/B]": 161,
+    "[/Br]": 62,
+    "[/C-1]": 229,
+    "[/C@@H1]": 45,
+    "[/C@@]": 105,
+    "[/C@H1]": 87,
+    "[/C@]": 76,
+    "[/C]": 20,
+    "[/C][=Branch1]": 428,
+    "[/C][=Branch2]": 450,
+    "[/C][=C]": 364,
+    "[/C][=N]": 454,
+    "[/C][Branch1]": 540,
+    "[/C][C]": 564,
+    "[/Cl]": 101,
+    "[/F]": 160,
+    "[/I]": 94,
+    "[/N+1]": 68,
+    "[/N-1]": 142,
+    "[/NH1]": 114,
+    "[/N]": 48,
+    "[/N][=C]": 418,
+    "[/N][=N]": 580,
+    "[/N][C]": 499,
+    "[/O-1]": 128,
+    "[/O]": 53,
+    "[/P+1]": 181,
+    "[/P@@]": 231,
+    "[/P]": 102,
+    "[/S+1]": 67,
+    "[/S-1]": 258,
+    "[/SH0]": 303,
+    "[/S]": 71,
+    "[/Se]": 149,
+    "[/Si]": 125,
+    "[10B]": 283,
+    "[11C@@H1]": 304,
+    "[11CH2]": 249,
+    "[11CH3]": 74,
+    "[11C]": 137,
+    "[123I-1]": 312,
+    "[123I]": 69,
+    "[123Te]": 186,
+    "[124I-1]": 325,
+    "[124I]": 133,
+    "[125I-1]": 171,
+    "[125I]": 70,
+    "[127I]": 301,
+    "[127Xe]": 311,
+    "[129Xe]": 292,
+    "[131Cs]": 275,
+    "[131I-1]": 329,
+    "[131I]": 86,
+    "[133Xe]": 322,
+    "[13CH1]": 198,
+    "[13CH3]": 277,
+    "[13C]": 166,
+    "[14C@@H1]": 318,
+    "[14C@@]": 302,
+    "[14C@H1]": 306,
+    "[14CH1]": 185,
+    "[14CH2]": 175,
+    "[14CH3]": 177,
+    "[14C]": 143,
+    "[15NH1]": 200,
+    "[15N]": 134,
+    "[17F]": 218,
+    "[18F-1]": 294,
+    "[18F]": 73,
+    "[18OH1]": 252,
+    "[19F]": 144,
+    "[211At]": 286,
+    "[223Ra]": 193,
+    "[22Na+1]": 191,
+    "[2H]": 88,
+    "[32P]": 307,
+    "[35S]": 182,
+    "[3H]": 103,
+    "[42K+1]": 188,
+    "[45Ca+2]": 178,
+    "[47Ca+2]": 201,
+    "[68Ga+3]": 273,
+    "[73Se]": 242,
+    "[75Se]": 327,
+    "[76Br]": 152,
+    "[81Kr]": 326,
+    "[82Rb+1]": 314,
+    "[82Rb]": 280,
+    "[85Sr+2]": 180,
+    "[89Sr+2]": 319,
+    "[=11C]": 132,
+    "[=13CH1]": 196,
+    "[=13C]": 199,
+    "[=14CH1]": 156,
+    "[=14C]": 194,
+    "[=18O]": 317,
+    "[=32PH1]": 257,
+    "[=32P]": 244,
+    "[=Al]": 295,
+    "[=As]": 145,
+    "[=B-1]": 159,
+    "[=B]": 176,
+    "[=Branch1]": 13,
+    "[=Branch1][=C]": 451,
+    "[=Branch1][=N]": 563,
+    "[=Branch1][=O]": 546,
+    "[=Branch1][C]": 335,
+    "[=Branch1][N+1]": 448,
+    "[=Branch1][N]": 508,
+    "[=Branch1][O]": 518,
+    "[=Branch2]": 14,
+    "[=Branch2][=C]": 472,
+    "[=Branch2][=N]": 584,
+    "[=Branch2][=O]": 549,
+    "[=Branch2][=Ring1]": 556,
+    "[=Branch2][C]": 383,
+    "[=Branch2][N]": 424,
+    "[=Branch2][O]": 507,
+    "[=Branch2][S]": 509,
+    "[=Branch3]": 592,
+    "[=C+1]": 267,
+    "[=CH0]": 593,
+    "[=C]": 7,
+    "[=C][=C]": 492,
+    "[=C][=O]": 516,
+    "[=C][Branch1]": 467,
+    "[=C][C]": 388,
+    "[=C][N]": 437,
+    "[=C][O]": 495,
+    "[=C][Ring1]": 482,
+    "[=C][S]": 456,
+    "[=Mg]": 310,
+    "[=N+1]": 32,
+    "[=N-1]": 63,
+    "[=NH1+1]": 232,
+    "[=NH2+1]": 243,
+    "[=N]": 29,
+    "[=N][=C]": 473,
+    "[=N][=N]": 579,
+    "[=N][=O]": 517,
+    "[=N][C]": 362,
+    "[=N][NH1]": 443,
+    "[=N][N]": 382,
+    "[=N][O]": 427,
+    "[=N][Ring1]": 513,
+    "[=N][S]": 501,
+    "[=O+1]": 116,
+    "[=OH0]": 296,
+    "[=OH1+1]": 594,
+    "[=O]": 18,
+    "[=O][/C]": 572,
+    "[=O][=C]": 411,
+    "[=O][=N]": 583,
+    "[=O][=O]": 425,
+    "[=O][C@@H1]": 414,
+    "[=O][C@@]": 544,
+    "[=O][C@H1]": 431,
+    "[=O][C@]": 514,
+    "[=O][C]": 361,
+    "[=O][NH1]": 434,
+    "[=O][N]": 345,
+    "[=O][O-1]": 391,
+    "[=O][O]": 370,
+    "[=O][S]": 571,
+    "[=P@@]": 107,
+    "[=P@]": 129,
+    "[=PH1]": 147,
+    "[=PH2]": 204,
+    "[=P]": 57,
+    "[=Ring1]": 34,
+    "[=Ring1][#Branch1]": 570,
+    "[=Ring1][#Branch2]": 576,
+    "[=Ring1][#C]": 562,
+    "[=Ring1][=Branch1]": 539,
+    "[=Ring1][=Branch2]": 481,
+    "[=Ring1][=C]": 582,
+    "[=Ring1][=N]": 545,
+    "[=Ring1][Branch1]": 442,
+    "[=Ring1][Branch2]": 568,
+    "[=Ring1][N]": 557,
+    "[=Ring1][O]": 526,
+    "[=Ring1][P]": 581,
+    "[=Ring1][S]": 565,
+    "[=Ring2]": 41,
+    "[=Ring2][Ring1]": 506,
+    "[=S+1]": 98,
+    "[=S-1]": 278,
+    "[=S@@]": 50,
+    "[=S@]": 164,
+    "[=SH0]": 269,
+    "[=SH1]": 150,
+    "[=S]": 26,
+    "[=S][N]": 455,
+    "[=Se+1]": 174,
+    "[=Se]": 162,
+    "[=Si]": 173,
+    "[=Te+1]": 179,
+    "[=Te]": 253,
+    "[Ag+1]": 190,
+    "[Ag-4]": 328,
+    "[Ag]": 146,
+    "[Al+3]": 223,
+    "[Al]": 225,
+    "[Ar]": 266,
+    "[As+1]": 189,
+    "[As-1]": 298,
+    "[AsH1]": 282,
+    "[AsH3]": 316,
+    "[As]": 84,
+    "[At]": 290,
+    "[Au-1]": 595,
+    "[B-1]": 100,
+    "[B@-1]": 206,
+    "[B@@-1]": 208,
+    "[BH0]": 165,
+    "[BH1-1]": 251,
+    "[BH2-1]": 202,
+    "[BH3-1]": 131,
+    "[B]": 65,
+    "[Ba+2]": 157,
+    "[Ba]": 291,
+    "[Be+2]": 246,
+    "[Bi+3]": 169,
+    "[Bi]": 155,
+    "[Br+2]": 254,
+    "[Br-1]": 52,
+    "[BrH0]": 596,
+    "[Br]": 21,
+    "[Br][=C]": 432,
+    "[Br][C]": 558,
+    "[Branch1]": 11,
+    "[Branch1][#Branch1]": 410,
+    "[Branch1][#Branch2]": 392,
+    "[Branch1][#C]": 399,
+    "[Branch1][=Branch1]": 401,
+    "[Branch1][=Branch2]": 358,
+    "[Branch1][=C]": 379,
+    "[Branch1][=N]": 405,
+    "[Branch1][=O]": 532,
+    "[Branch1][Branch1]": 476,
+    "[Branch1][Branch2]": 394,
+    "[Branch1][C]": 336,
+    "[Branch1][N]": 378,
+    "[Branch1][O]": 398,
+    "[Branch1][P]": 403,
+    "[Branch1][Ring1]": 396,
+    "[Branch1][Ring2]": 429,
+    "[Branch1][S]": 402,
+    "[Branch2]": 8,
+    "[Branch2][#Branch1]": 551,
+    "[Branch2][=Branch1]": 446,
+    "[Branch2][=C]": 464,
+    "[Branch2][=N]": 560,
+    "[Branch2][=O]": 522,
+    "[Branch2][Branch1]": 384,
+    "[Branch2][C]": 409,
+    "[Branch2][N]": 440,
+    "[Branch2][O]": 478,
+    "[Branch2][Ring1]": 340,
+    "[Branch2][Ring2]": 350,
+    "[Branch2][S]": 535,
+    "[Branch3]": 597,
+    "[C+1]": 106,
+    "[C-1]": 119,
+    "[C@@H1]": 27,
+    "[C@@H1][Branch1]": 371,
+    "[C@@H1][C@@H1]": 547,
+    "[C@@H1][C@H1]": 586,
+    "[C@@H1][C]": 447,
+    "[C@@H1][O]": 493,
+    "[C@@H1][Ring1]": 524,
+    "[C@@]": 28,
+    "[C@@][Branch1]": 413,
+    "[C@H1]": 33,
+    "[C@H1][Branch1]": 366,
+    "[C@H1][Branch2]": 573,
+    "[C@H1][C@@H1]": 505,
+    "[C@H1][C@H1]": 575,
+    "[C@H1][C]": 445,
+    "[C@H1][O]": 530,
+    "[C@H1][Ring1]": 470,
+    "[C@]": 35,
+    "[C@][Branch1]": 412,
+    "[CH0]": 227,
+    "[CH1-1]": 138,
+    "[CH1]": 323,
+    "[CH2-1]": 598,
+    "[CH2]": 224,
+    "[CH3]": 226,
+    "[C]": 5,
+    "[C][#C]": 500,
+    "[C][#N]": 536,
+    "[C][=Branch1]": 377,
+    "[C][=Branch2]": 488,
+    "[C][=C]": 333,
+    "[C][=N+1]": 543,
+    "[C][=N]": 346,
+    "[C][=O]": 511,
+    "[C][=Ring1]": 395,
+    "[C][=Ring2]": 471,
+    "[C][Branch1]": 342,
+    "[C][Branch2]": 423,
+    "[C][C@@H1]": 386,
+    "[C][C@@]": 480,
+    "[C][C@H1]": 385,
+    "[C][C@]": 462,
+    "[C][C]": 334,
+    "[C][Cl]": 574,
+    "[C][F]": 578,
+    "[C][NH1]": 449,
+    "[C][N]": 352,
+    "[C][O]": 348,
+    "[C][Ring1]": 351,
+    "[C][Ring2]": 588,
+    "[C][S]": 397,
+    "[Ca+2]": 122,
+    "[CaH2]": 288,
+    "[Ca]": 234,
+    "[Cl+1]": 139,
+    "[Cl+2]": 240,
+    "[Cl+3]": 61,
+    "[Cl-1]": 39,
+    "[ClH0]": 126,
+    "[ClH1+1]": 599,
+    "[ClH2+1]": 600,
+    "[Cl]": 24,
+    "[Cl][=C]": 374,
+    "[Cl][=N]": 494,
+    "[Cl][C]": 416,
+    "[Cl][Cl]": 484,
+    "[Cl][N]": 512,
+    "[Cl][O]": 587,
+    "[Cl][S]": 555,
+    "[Co-2]": 601,
+    "[Co-4]": 602,
+    "[Cr]": 603,
+    "[Cs+1]": 85,
+    "[Cs]": 284,
+    "[Cu-2]": 604,
+    "[Cu-3]": 605,
+    "[Cu-5]": 606,
+    "[Cu]": 607,
+    "[F-1]": 111,
+    "[F]": 16,
+    "[F][=C]": 376,
+    "[F][=N]": 537,
+    "[F][C]": 404,
+    "[F][F]": 360,
+    "[F][N]": 585,
+    "[F][O]": 577,
+    "[Fe+2]": 608,
+    "[Fe+3]": 609,
+    "[Fe-2]": 610,
+    "[Fe-3]": 611,
+    "[Gd+3]": 612,
+    "[Ge]": 613,
+    "[H+1]": 151,
+    "[H-1]": 285,
+    "[H]": 614,
+    "[He]": 115,
+    "[Hg]": 615,
+    "[I+1]": 153,
+    "[I+2]": 308,
+    "[I+3]": 309,
+    "[I-1]": 66,
+    "[IH0]": 183,
+    "[I]": 42,
+    "[Ir-3]": 616,
+    "[K+1]": 83,
+    "[KH1]": 228,
+    "[K]": 313,
+    "[Kr]": 315,
+    "[Li+1]": 110,
+    "[LiH1]": 233,
+    "[Li]": 276,
+    "[Mg+1]": 289,
+    "[Mg+2]": 75,
+    "[MgH2]": 236,
+    "[Mg]": 235,
+    "[Mn+1]": 617,
+    "[Mn+2]": 618,
+    "[Mo]": 619,
+    "[N+1]": 30,
+    "[N+1][=C]": 510,
+    "[N+1][Branch1]": 497,
+    "[N-1]": 37,
+    "[N@+1]": 97,
+    "[N@@+1]": 109,
+    "[N@@]": 140,
+    "[N@H1+1]": 297,
+    "[N@]": 91,
+    "[NH0]": 238,
+    "[NH1+1]": 216,
+    "[NH1-1]": 124,
+    "[NH1]": 38,
+    "[NH1][C]": 406,
+    "[NH1][Ring1]": 566,
+    "[NH2+1]": 215,
+    "[NH2-1]": 620,
+    "[NH2]": 621,
+    "[NH3+1]": 293,
+    "[NH4+1]": 260,
+    "[N]": 10,
+    "[N][#C]": 422,
+    "[N][=C]": 343,
+    "[N][=N]": 389,
+    "[N][=O]": 417,
+    "[N][=Ring1]": 420,
+    "[N][=Ring2]": 457,
+    "[N][Branch1]": 347,
+    "[N][Branch2]": 498,
+    "[N][C@@H1]": 486,
+    "[N][C@H1]": 483,
+    "[N][C]": 339,
+    "[N][N]": 430,
+    "[N][O]": 468,
+    "[N][Ring1]": 415,
+    "[N][S]": 408,
+    "[Na+1]": 40,
+    "[NaH1]": 221,
+    "[Na]": 268,
+    "[Ni-2]": 622,
+    "[Ni-4]": 623,
+    "[O+1]": 136,
+    "[O-1]": 31,
+    "[O-2]": 270,
+    "[OH0]": 96,
+    "[OH1+1]": 220,
+    "[OH1-1]": 154,
+    "[OH1]": 212,
+    "[OH2+1]": 624,
+    "[OH3+1]": 274,
+    "[O]": 6,
+    "[O][=C]": 349,
+    "[O][=N+1]": 503,
+    "[O][=N]": 474,
+    "[O][=O]": 521,
+    "[O][=S]": 460,
+    "[O][Branch1]": 559,
+    "[O][C@@H1]": 390,
+    "[O][C@@]": 553,
+    "[O][C@H1]": 387,
+    "[O][C@]": 519,
+    "[O][C]": 341,
+    "[O][N]": 441,
+    "[O][O]": 426,
+    "[O][P]": 489,
+    "[O][Ring1]": 458,
+    "[O][S]": 487,
+    "[P+1]": 112,
+    "[P-1]": 135,
+    "[P@+1]": 163,
+    "[P@@+1]": 237,
+    "[P@@]": 78,
+    "[P@]": 72,
+    "[PH1]": 127,
+    "[PH2+1]": 305,
+    "[P]": 25,
+    "[P][=C]": 504,
+    "[P][=O]": 552,
+    "[P][C]": 419,
+    "[P][N]": 477,
+    "[P][O]": 548,
+    "[P][S]": 533,
+    "[Pd-2]": 625,
+    "[Pt-2]": 626,
+    "[Pt]": 627,
+    "[Ra]": 187,
+    "[Rb+1]": 207,
+    "[Rb]": 321,
+    "[Rh+2]": 628,
+    "[Ring1]": 9,
+    "[Ring1][#Branch1]": 344,
+    "[Ring1][#Branch2]": 353,
+    "[Ring1][#C]": 369,
+    "[Ring1][=Branch1]": 338,
+    "[Ring1][=Branch2]": 356,
+    "[Ring1][=C]": 363,
+    "[Ring1][=N]": 367,
+    "[Ring1][=O]": 550,
+    "[Ring1][Branch1]": 355,
+    "[Ring1][Branch2]": 357,
+    "[Ring1][C]": 373,
+    "[Ring1][N]": 368,
+    "[Ring1][O]": 354,
+    "[Ring1][P]": 372,
+    "[Ring1][Ring1]": 421,
+    "[Ring1][Ring2]": 439,
+    "[Ring1][S]": 365,
+    "[Ring2]": 12,
+    "[Ring2][=C]": 381,
+    "[Ring2][=N]": 459,
+    "[Ring2][=O]": 525,
+    "[Ring2][Branch1]": 520,
+    "[Ring2][C]": 393,
+    "[Ring2][N]": 433,
+    "[Ring2][O]": 466,
+    "[Ring2][Ring1]": 337,
+    "[Ring2][Ring2]": 359,
+    "[Ring2][S]": 502,
+    "[Ring3]": 629,
+    "[S+1]": 44,
+    "[S-1]": 60,
+    "[S-2]": 172,
+    "[S@+1]": 56,
+    "[S@@+1]": 95,
+    "[S@@]": 51,
+    "[S@]": 184,
+    "[SH0]": 170,
+    "[SH1+1]": 263,
+    "[SH1-1]": 205,
+    "[SH1]": 264,
+    "[SH2]": 299,
+    "[S]": 17,
+    "[S][=C]": 490,
+    "[S][=N]": 523,
+    "[S][=O]": 554,
+    "[S][Branch1]": 444,
+    "[S][C]": 375,
+    "[S][N]": 452,
+    "[S][O]": 527,
+    "[S][Ring1]": 485,
+    "[S][S]": 463,
+    "[Sb]": 211,
+    "[Se+1]": 230,
+    "[SeH1]": 130,
+    "[SeH2]": 330,
+    "[Se]": 36,
+    "[Si-1]": 331,
+    "[Si@]": 287,
+    "[SiH1-1]": 192,
+    "[SiH1]": 104,
+    "[SiH2]": 300,
+    "[SiH3-1]": 219,
+    "[SiH4]": 271,
+    "[Si]": 49,
+    "[Sn]": 630,
+    "[Sr+2]": 158,
+    "[Sr]": 332,
+    "[TeH1]": 279,
+    "[TeH2]": 239,
+    "[Te]": 81,
+    "[Xe]": 320,
+    "[Yb]": 209,
+    "[Zn+1]": 213,
+    "[Zn+2]": 148,
+    "[Zn-2]": 261,
+    "[Zn]": 82,
+    "[\\11CH3]": 203,
+    "[\\123I]": 214,
+    "[\\125I]": 79,
+    "[\\3H]": 250,
+    "[\\B]": 217,
+    "[\\Br]": 89,
+    "[\\C-1]": 265,
+    "[\\C@@H1]": 46,
+    "[\\C@@]": 80,
+    "[\\C@H1]": 77,
+    "[\\C@]": 99,
+    "[\\CH1-1]": 241,
+    "[\\C]": 43,
+    "[\\C][=C]": 436,
+    "[\\C][C]": 567,
+    "[\\Cl]": 58,
+    "[\\F]": 118,
+    "[\\I]": 59,
+    "[\\N+1]": 113,
+    "[\\N-1]": 222,
+    "[\\NH1]": 108,
+    "[\\N]": 55,
+    "[\\N][C]": 541,
+    "[\\O-1]": 64,
+    "[\\O]": 47,
+    "[\\P]": 93,
+    "[\\S+1]": 121,
+    "[\\S-1]": 255,
+    "[\\S]": 54,
+    "[\\Se]": 168,
+    "[\\Si]": 195
+}

ape_tokenizer_metadata.json

@@ -0,0 +1,102 @@
+{
+  "dataset_name": "data/pretrain/chembl36_selfies",
+  "selfies_column": "selfies",
+  "train_split": "train",
+  "validation_split": "valid",
+  "use_validation_split": true,
+  "representation": "SELFIES",
+  "expected_input": "SELFIES strings only. Convert SMILES before inference using a helper such as smiles_to_selfies().",
+  "tokenizer_vocab_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+  "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json",
+  "backend": "cuda",
+  "platform": "Linux-6.8.0-110-generic-x86_64-with-glibc2.35",
+  "torch_version": "2.11.0+cu130",
+  "transformers_version": "5.8.1",
+  "vocab_size": 631,
+  "special_ids": {
+    "pad_token": 1,
+    "bos_token": 0,
+    "eos_token": 2,
+    "unk_token": 3,
+    "mask_token": 4
+  },
+  "num_parameters": 34149495,
+  "tokenizer_stats": {
+    "sample_size": 1000.0,
+    "mean_len": 25.555,
+    "p50_len": 25.0,
+    "p95_len": 41.0,
+    "p99_len": 51.0,
+    "max_len": 62.0,
+    "truncation_rate": 0.0,
+    "unk_rate": 0.0,
+    "empty_sequence_rate": 0.0,
+    "mostly_unknown_rate": 0.0
+  },
+  "final_eval_metrics": {
+    "eval_loss": 0.3744058609008789,
+    "eval_masked_accuracy": 0.8805643239064613,
+    "eval_runtime": 2.8656,
+    "eval_samples_per_second": 1429.379,
+    "eval_steps_per_second": 5.584,
+    "epoch": 3.066,
+    "eval_perplexity": 1.454127204085947
+  },
+  "trainer_state_summary": {
+    "best_global_step": 30000,
+    "best_metric": 0.37694016098976135,
+    "best_model_checkpoint": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_4e-4/checkpoint-30000",
+    "global_step": 30000
+  },
+  "args": {
+    "output_dir": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_4e-4",
+    "tokenizer_vocab_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+    "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json",
+    "dataset_name": "data/pretrain/chembl36_selfies",
+    "selfies_column": "selfies",
+    "train_split": "train",
+    "validation_split": "valid",
+    "use_validation_split": true,
+    "data_dir": null,
+    "data_files": null,
+    "eval_size": 4096,
+    "shuffle_buffer_size": 100000,
+    "seed": 42,
+    "val_split_mod": 100,
+    "val_split_bucket": 0,
+    "tokenizer_validation_samples": 1000,
+    "unk_rate_threshold": 0.001,
+    "truncation_warn_threshold": 0.05,
+    "model_size": "small",
+    "max_seq_length": 128,
+    "mlm_probability": 0.15,
+    "masking_strategy": "standard",
+    "span_p_geom": 0.4,
+    "span_max_length": 6,
+    "heteroatom_start_weight": 2.0,
+    "max_steps": 30000,
+    "per_device_train_batch_size": 256,
+    "per_device_eval_batch_size": 256,
+    "gradient_accumulation_steps": 1,
+    "learning_rate": 0.0004,
+    "weight_decay": 0.01,
+    "warmup_steps": 1500,
+    "max_grad_norm": 1.0,
+    "load_best_model_at_end": true,
+    "metric_for_best_model": "eval_loss",
+    "greater_is_better": false,
+    "logging_steps": 100,
+    "eval_steps": 5000,
+    "save_steps": 5000,
+    "save_total_limit": 2,
+    "device_backend": "cuda",
+    "bf16": true,
+    "fp16": false,
+    "num_workers": 4,
+    "max_eval_batches": 16,
+    "report_to": "tensorboard",
+    "compute_masked_accuracy": true,
+    "debug": false,
+    "hf_login": false
+  }
+}

collator_config.json

@@ -0,0 +1,5 @@
+{
+  "_note": "Collator parameters used during pre-training. Change masking_strategy to 'standard', 'span', or 'hetero_span' and adjust the corresponding parameters when fine-tuning.",
+  "masking_strategy": "standard",
+  "mlm_probability": 0.15
+}

config.json

@@ -9,7 +9,7 @@
   "classifier_bias": false,
   "classifier_dropout": 0.0,
   "classifier_pooling": "mean",
-  "cls_token_id": 50281,
+  "cls_token_id": 0,
   "decoder_bias": true,
   "deterministic_flash_attn": false,
   "dtype": "float32",
@@ -34,6 +34,7 @@
     "sliding_attention"
   ],
   "local_attention": 128,
+  "mask_token_id": 4,
   "max_position_embeddings": 128,
   "mlp_bias": false,
   "mlp_dropout": 0.0,
@@ -54,11 +55,12 @@
       "rope_type": "default"
     }
   },
-  "sep_token_id": 50282,
+  "sep_token_id": 2,
   "sparse_pred_ignore_index": -100,
   "sparse_prediction": false,
   "tie_word_embeddings": true,
   "transformers_version": "5.8.1",
+  "unk_token_id": 3,
   "use_cache": false,
-  "vocab_size": 2439
+  "vocab_size": 631
 }

eval_results.json

@@ -0,0 +1,9 @@
+{
+    "epoch": 3.066,
+    "eval_loss": 0.3744058609008789,
+    "eval_masked_accuracy": 0.8805643239064613,
+    "eval_perplexity": 1.454127204085947,
+    "eval_runtime": 2.8656,
+    "eval_samples_per_second": 1429.379,
+    "eval_steps_per_second": 5.584
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eaeb9ab56962d39645f5a39f85f65d1f759c7b6904fa5b301bf3d07803a1c934
-size 140313324
+oid sha256:9520e881983046b32a4b2f8c65665e8264fe9a4887935bf89cd8ee4842d057cd
+size 136603308

run_args.json

@@ -1,8 +1,8 @@
 {
-  "output_dir": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_1e-4",
-  "tokenizer_vocab_path": "tokenizer/chembl36_selfies_2m_benchmark_covered_ape_tokenizer.json",
-  "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_benchmark_covered_ape_tokenizer.metadata.json",
-  "dataset_name": "data/pretrain/chembl36_selfies_tokenized",
+  "output_dir": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_4e-4",
+  "tokenizer_vocab_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+  "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json",
+  "dataset_name": "data/pretrain/chembl36_selfies",
   "selfies_column": "selfies",
   "train_split": "train",
   "validation_split": "valid",
@@ -28,7 +28,7 @@
   "per_device_train_batch_size": 256,
   "per_device_eval_batch_size": 256,
   "gradient_accumulation_steps": 1,
-  "learning_rate": 0.0001,
+  "learning_rate": 0.0004,
   "weight_decay": 0.01,
   "warmup_steps": 1500,
   "max_grad_norm": 1.0,
@@ -43,7 +43,7 @@
   "bf16": true,
   "fp16": false,
   "num_workers": 4,
-  "max_eval_batches": 256,
+  "max_eval_batches": 16,
   "report_to": "tensorboard",
   "compute_masked_accuracy": true,
   "debug": false,

selfies_vocab.json

@@ -0,0 +1,633 @@
+{
+    "</s>": 2,
+    "<mask>": 4,
+    "<pad>": 1,
+    "<s>": 0,
+    "<unk>": 3,
+    "[#11C]": 256,
+    "[#14C]": 247,
+    "[#Branch1]": 22,
+    "[#Branch1][=C]": 479,
+    "[#Branch1][=O]": 542,
+    "[#Branch1][C]": 400,
+    "[#Branch1][N]": 453,
+    "[#Branch1][O]": 491,
+    "[#Branch1][S]": 528,
+    "[#Branch2]": 15,
+    "[#Branch2][=C]": 469,
+    "[#Branch2][=N]": 561,
+    "[#Branch2][=O]": 531,
+    "[#Branch2][C]": 380,
+    "[#Branch2][N]": 438,
+    "[#Branch2][O]": 461,
+    "[#Branch2][S]": 529,
+    "[#C-1]": 141,
+    "[#CH0]": 248,
+    "[#C]": 23,
+    "[#C][=C]": 496,
+    "[#C][=N]": 569,
+    "[#C][=O]": 515,
+    "[#C][C]": 407,
+    "[#C][N]": 465,
+    "[#C][O]": 538,
+    "[#C][S]": 534,
+    "[#N+1]": 123,
+    "[#N]": 19,
+    "[#N][=C]": 435,
+    "[#N][C]": 475,
+    "[#O+1]": 589,
+    "[#S]": 262,
+    "[-/Ring1]": 92,
+    "[-/Ring2]": 90,
+    "[-/Ring3]": 590,
+    "[-\\Ring1]": 167,
+    "[-\\Ring2]": 120,
+    "[-\\Ring3]": 591,
+    "[/123I]": 245,
+    "[/125I]": 117,
+    "[/131I]": 210,
+    "[/13CH1]": 272,
+    "[/13C]": 197,
+    "[/14CH1]": 324,
+    "[/14C]": 281,
+    "[/As]": 259,
+    "[/B]": 161,
+    "[/Br]": 62,
+    "[/C-1]": 229,
+    "[/C@@H1]": 45,
+    "[/C@@]": 105,
+    "[/C@H1]": 87,
+    "[/C@]": 76,
+    "[/C]": 20,
+    "[/C][=Branch1]": 428,
+    "[/C][=Branch2]": 450,
+    "[/C][=C]": 364,
+    "[/C][=N]": 454,
+    "[/C][Branch1]": 540,
+    "[/C][C]": 564,
+    "[/Cl]": 101,
+    "[/F]": 160,
+    "[/I]": 94,
+    "[/N+1]": 68,
+    "[/N-1]": 142,
+    "[/NH1]": 114,
+    "[/N]": 48,
+    "[/N][=C]": 418,
+    "[/N][=N]": 580,
+    "[/N][C]": 499,
+    "[/O-1]": 128,
+    "[/O]": 53,
+    "[/P+1]": 181,
+    "[/P@@]": 231,
+    "[/P]": 102,
+    "[/S+1]": 67,
+    "[/S-1]": 258,
+    "[/SH0]": 303,
+    "[/S]": 71,
+    "[/Se]": 149,
+    "[/Si]": 125,
+    "[10B]": 283,
+    "[11C@@H1]": 304,
+    "[11CH2]": 249,
+    "[11CH3]": 74,
+    "[11C]": 137,
+    "[123I-1]": 312,
+    "[123I]": 69,
+    "[123Te]": 186,
+    "[124I-1]": 325,
+    "[124I]": 133,
+    "[125I-1]": 171,
+    "[125I]": 70,
+    "[127I]": 301,
+    "[127Xe]": 311,
+    "[129Xe]": 292,
+    "[131Cs]": 275,
+    "[131I-1]": 329,
+    "[131I]": 86,
+    "[133Xe]": 322,
+    "[13CH1]": 198,
+    "[13CH3]": 277,
+    "[13C]": 166,
+    "[14C@@H1]": 318,
+    "[14C@@]": 302,
+    "[14C@H1]": 306,
+    "[14CH1]": 185,
+    "[14CH2]": 175,
+    "[14CH3]": 177,
+    "[14C]": 143,
+    "[15NH1]": 200,
+    "[15N]": 134,
+    "[17F]": 218,
+    "[18F-1]": 294,
+    "[18F]": 73,
+    "[18OH1]": 252,
+    "[19F]": 144,
+    "[211At]": 286,
+    "[223Ra]": 193,
+    "[22Na+1]": 191,
+    "[2H]": 88,
+    "[32P]": 307,
+    "[35S]": 182,
+    "[3H]": 103,
+    "[42K+1]": 188,
+    "[45Ca+2]": 178,
+    "[47Ca+2]": 201,
+    "[68Ga+3]": 273,
+    "[73Se]": 242,
+    "[75Se]": 327,
+    "[76Br]": 152,
+    "[81Kr]": 326,
+    "[82Rb+1]": 314,
+    "[82Rb]": 280,
+    "[85Sr+2]": 180,
+    "[89Sr+2]": 319,
+    "[=11C]": 132,
+    "[=13CH1]": 196,
+    "[=13C]": 199,
+    "[=14CH1]": 156,
+    "[=14C]": 194,
+    "[=18O]": 317,
+    "[=32PH1]": 257,
+    "[=32P]": 244,
+    "[=Al]": 295,
+    "[=As]": 145,
+    "[=B-1]": 159,
+    "[=B]": 176,
+    "[=Branch1]": 13,
+    "[=Branch1][=C]": 451,
+    "[=Branch1][=N]": 563,
+    "[=Branch1][=O]": 546,
+    "[=Branch1][C]": 335,
+    "[=Branch1][N+1]": 448,
+    "[=Branch1][N]": 508,
+    "[=Branch1][O]": 518,
+    "[=Branch2]": 14,
+    "[=Branch2][=C]": 472,
+    "[=Branch2][=N]": 584,
+    "[=Branch2][=O]": 549,
+    "[=Branch2][=Ring1]": 556,
+    "[=Branch2][C]": 383,
+    "[=Branch2][N]": 424,
+    "[=Branch2][O]": 507,
+    "[=Branch2][S]": 509,
+    "[=Branch3]": 592,
+    "[=C+1]": 267,
+    "[=CH0]": 593,
+    "[=C]": 7,
+    "[=C][=C]": 492,
+    "[=C][=O]": 516,
+    "[=C][Branch1]": 467,
+    "[=C][C]": 388,
+    "[=C][N]": 437,
+    "[=C][O]": 495,
+    "[=C][Ring1]": 482,
+    "[=C][S]": 456,
+    "[=Mg]": 310,
+    "[=N+1]": 32,
+    "[=N-1]": 63,
+    "[=NH1+1]": 232,
+    "[=NH2+1]": 243,
+    "[=N]": 29,
+    "[=N][=C]": 473,
+    "[=N][=N]": 579,
+    "[=N][=O]": 517,
+    "[=N][C]": 362,
+    "[=N][NH1]": 443,
+    "[=N][N]": 382,
+    "[=N][O]": 427,
+    "[=N][Ring1]": 513,
+    "[=N][S]": 501,
+    "[=O+1]": 116,
+    "[=OH0]": 296,
+    "[=OH1+1]": 594,
+    "[=O]": 18,
+    "[=O][/C]": 572,
+    "[=O][=C]": 411,
+    "[=O][=N]": 583,
+    "[=O][=O]": 425,
+    "[=O][C@@H1]": 414,
+    "[=O][C@@]": 544,
+    "[=O][C@H1]": 431,
+    "[=O][C@]": 514,
+    "[=O][C]": 361,
+    "[=O][NH1]": 434,
+    "[=O][N]": 345,
+    "[=O][O-1]": 391,
+    "[=O][O]": 370,
+    "[=O][S]": 571,
+    "[=P@@]": 107,
+    "[=P@]": 129,
+    "[=PH1]": 147,
+    "[=PH2]": 204,
+    "[=P]": 57,
+    "[=Ring1]": 34,
+    "[=Ring1][#Branch1]": 570,
+    "[=Ring1][#Branch2]": 576,
+    "[=Ring1][#C]": 562,
+    "[=Ring1][=Branch1]": 539,
+    "[=Ring1][=Branch2]": 481,
+    "[=Ring1][=C]": 582,
+    "[=Ring1][=N]": 545,
+    "[=Ring1][Branch1]": 442,
+    "[=Ring1][Branch2]": 568,
+    "[=Ring1][N]": 557,
+    "[=Ring1][O]": 526,
+    "[=Ring1][P]": 581,
+    "[=Ring1][S]": 565,
+    "[=Ring2]": 41,
+    "[=Ring2][Ring1]": 506,
+    "[=S+1]": 98,
+    "[=S-1]": 278,
+    "[=S@@]": 50,
+    "[=S@]": 164,
+    "[=SH0]": 269,
+    "[=SH1]": 150,
+    "[=S]": 26,
+    "[=S][N]": 455,
+    "[=Se+1]": 174,
+    "[=Se]": 162,
+    "[=Si]": 173,
+    "[=Te+1]": 179,
+    "[=Te]": 253,
+    "[Ag+1]": 190,
+    "[Ag-4]": 328,
+    "[Ag]": 146,
+    "[Al+3]": 223,
+    "[Al]": 225,
+    "[Ar]": 266,
+    "[As+1]": 189,
+    "[As-1]": 298,
+    "[AsH1]": 282,
+    "[AsH3]": 316,
+    "[As]": 84,
+    "[At]": 290,
+    "[Au-1]": 595,
+    "[B-1]": 100,
+    "[B@-1]": 206,
+    "[B@@-1]": 208,
+    "[BH0]": 165,
+    "[BH1-1]": 251,
+    "[BH2-1]": 202,
+    "[BH3-1]": 131,
+    "[B]": 65,
+    "[Ba+2]": 157,
+    "[Ba]": 291,
+    "[Be+2]": 246,
+    "[Bi+3]": 169,
+    "[Bi]": 155,
+    "[Br+2]": 254,
+    "[Br-1]": 52,
+    "[BrH0]": 596,
+    "[Br]": 21,
+    "[Br][=C]": 432,
+    "[Br][C]": 558,
+    "[Branch1]": 11,
+    "[Branch1][#Branch1]": 410,
+    "[Branch1][#Branch2]": 392,
+    "[Branch1][#C]": 399,
+    "[Branch1][=Branch1]": 401,
+    "[Branch1][=Branch2]": 358,
+    "[Branch1][=C]": 379,
+    "[Branch1][=N]": 405,
+    "[Branch1][=O]": 532,
+    "[Branch1][Branch1]": 476,
+    "[Branch1][Branch2]": 394,
+    "[Branch1][C]": 336,
+    "[Branch1][N]": 378,
+    "[Branch1][O]": 398,
+    "[Branch1][P]": 403,
+    "[Branch1][Ring1]": 396,
+    "[Branch1][Ring2]": 429,
+    "[Branch1][S]": 402,
+    "[Branch2]": 8,
+    "[Branch2][#Branch1]": 551,
+    "[Branch2][=Branch1]": 446,
+    "[Branch2][=C]": 464,
+    "[Branch2][=N]": 560,
+    "[Branch2][=O]": 522,
+    "[Branch2][Branch1]": 384,
+    "[Branch2][C]": 409,
+    "[Branch2][N]": 440,
+    "[Branch2][O]": 478,
+    "[Branch2][Ring1]": 340,
+    "[Branch2][Ring2]": 350,
+    "[Branch2][S]": 535,
+    "[Branch3]": 597,
+    "[C+1]": 106,
+    "[C-1]": 119,
+    "[C@@H1]": 27,
+    "[C@@H1][Branch1]": 371,
+    "[C@@H1][C@@H1]": 547,
+    "[C@@H1][C@H1]": 586,
+    "[C@@H1][C]": 447,
+    "[C@@H1][O]": 493,
+    "[C@@H1][Ring1]": 524,
+    "[C@@]": 28,
+    "[C@@][Branch1]": 413,
+    "[C@H1]": 33,
+    "[C@H1][Branch1]": 366,
+    "[C@H1][Branch2]": 573,
+    "[C@H1][C@@H1]": 505,
+    "[C@H1][C@H1]": 575,
+    "[C@H1][C]": 445,
+    "[C@H1][O]": 530,
+    "[C@H1][Ring1]": 470,
+    "[C@]": 35,
+    "[C@][Branch1]": 412,
+    "[CH0]": 227,
+    "[CH1-1]": 138,
+    "[CH1]": 323,
+    "[CH2-1]": 598,
+    "[CH2]": 224,
+    "[CH3]": 226,
+    "[C]": 5,
+    "[C][#C]": 500,
+    "[C][#N]": 536,
+    "[C][=Branch1]": 377,
+    "[C][=Branch2]": 488,
+    "[C][=C]": 333,
+    "[C][=N+1]": 543,
+    "[C][=N]": 346,
+    "[C][=O]": 511,
+    "[C][=Ring1]": 395,
+    "[C][=Ring2]": 471,
+    "[C][Branch1]": 342,
+    "[C][Branch2]": 423,
+    "[C][C@@H1]": 386,
+    "[C][C@@]": 480,
+    "[C][C@H1]": 385,
+    "[C][C@]": 462,
+    "[C][C]": 334,
+    "[C][Cl]": 574,
+    "[C][F]": 578,
+    "[C][NH1]": 449,
+    "[C][N]": 352,
+    "[C][O]": 348,
+    "[C][Ring1]": 351,
+    "[C][Ring2]": 588,
+    "[C][S]": 397,
+    "[Ca+2]": 122,
+    "[CaH2]": 288,
+    "[Ca]": 234,
+    "[Cl+1]": 139,
+    "[Cl+2]": 240,
+    "[Cl+3]": 61,
+    "[Cl-1]": 39,
+    "[ClH0]": 126,
+    "[ClH1+1]": 599,
+    "[ClH2+1]": 600,
+    "[Cl]": 24,
+    "[Cl][=C]": 374,
+    "[Cl][=N]": 494,
+    "[Cl][C]": 416,
+    "[Cl][Cl]": 484,
+    "[Cl][N]": 512,
+    "[Cl][O]": 587,
+    "[Cl][S]": 555,
+    "[Co-2]": 601,
+    "[Co-4]": 602,
+    "[Cr]": 603,
+    "[Cs+1]": 85,
+    "[Cs]": 284,
+    "[Cu-2]": 604,
+    "[Cu-3]": 605,
+    "[Cu-5]": 606,
+    "[Cu]": 607,
+    "[F-1]": 111,
+    "[F]": 16,
+    "[F][=C]": 376,
+    "[F][=N]": 537,
+    "[F][C]": 404,
+    "[F][F]": 360,
+    "[F][N]": 585,
+    "[F][O]": 577,
+    "[Fe+2]": 608,
+    "[Fe+3]": 609,
+    "[Fe-2]": 610,
+    "[Fe-3]": 611,
+    "[Gd+3]": 612,
+    "[Ge]": 613,
+    "[H+1]": 151,
+    "[H-1]": 285,
+    "[H]": 614,
+    "[He]": 115,
+    "[Hg]": 615,
+    "[I+1]": 153,
+    "[I+2]": 308,
+    "[I+3]": 309,
+    "[I-1]": 66,
+    "[IH0]": 183,
+    "[I]": 42,
+    "[Ir-3]": 616,
+    "[K+1]": 83,
+    "[KH1]": 228,
+    "[K]": 313,
+    "[Kr]": 315,
+    "[Li+1]": 110,
+    "[LiH1]": 233,
+    "[Li]": 276,
+    "[Mg+1]": 289,
+    "[Mg+2]": 75,
+    "[MgH2]": 236,
+    "[Mg]": 235,
+    "[Mn+1]": 617,
+    "[Mn+2]": 618,
+    "[Mo]": 619,
+    "[N+1]": 30,
+    "[N+1][=C]": 510,
+    "[N+1][Branch1]": 497,
+    "[N-1]": 37,
+    "[N@+1]": 97,
+    "[N@@+1]": 109,
+    "[N@@]": 140,
+    "[N@H1+1]": 297,
+    "[N@]": 91,
+    "[NH0]": 238,
+    "[NH1+1]": 216,
+    "[NH1-1]": 124,
+    "[NH1]": 38,
+    "[NH1][C]": 406,
+    "[NH1][Ring1]": 566,
+    "[NH2+1]": 215,
+    "[NH2-1]": 620,
+    "[NH2]": 621,
+    "[NH3+1]": 293,
+    "[NH4+1]": 260,
+    "[N]": 10,
+    "[N][#C]": 422,
+    "[N][=C]": 343,
+    "[N][=N]": 389,
+    "[N][=O]": 417,
+    "[N][=Ring1]": 420,
+    "[N][=Ring2]": 457,
+    "[N][Branch1]": 347,
+    "[N][Branch2]": 498,
+    "[N][C@@H1]": 486,
+    "[N][C@H1]": 483,
+    "[N][C]": 339,
+    "[N][N]": 430,
+    "[N][O]": 468,
+    "[N][Ring1]": 415,
+    "[N][S]": 408,
+    "[Na+1]": 40,
+    "[NaH1]": 221,
+    "[Na]": 268,
+    "[Ni-2]": 622,
+    "[Ni-4]": 623,
+    "[O+1]": 136,
+    "[O-1]": 31,
+    "[O-2]": 270,
+    "[OH0]": 96,
+    "[OH1+1]": 220,
+    "[OH1-1]": 154,
+    "[OH1]": 212,
+    "[OH2+1]": 624,
+    "[OH3+1]": 274,
+    "[O]": 6,
+    "[O][=C]": 349,
+    "[O][=N+1]": 503,
+    "[O][=N]": 474,
+    "[O][=O]": 521,
+    "[O][=S]": 460,
+    "[O][Branch1]": 559,
+    "[O][C@@H1]": 390,
+    "[O][C@@]": 553,
+    "[O][C@H1]": 387,
+    "[O][C@]": 519,
+    "[O][C]": 341,
+    "[O][N]": 441,
+    "[O][O]": 426,
+    "[O][P]": 489,
+    "[O][Ring1]": 458,
+    "[O][S]": 487,
+    "[P+1]": 112,
+    "[P-1]": 135,
+    "[P@+1]": 163,
+    "[P@@+1]": 237,
+    "[P@@]": 78,
+    "[P@]": 72,
+    "[PH1]": 127,
+    "[PH2+1]": 305,
+    "[P]": 25,
+    "[P][=C]": 504,
+    "[P][=O]": 552,
+    "[P][C]": 419,
+    "[P][N]": 477,
+    "[P][O]": 548,
+    "[P][S]": 533,
+    "[Pd-2]": 625,
+    "[Pt-2]": 626,
+    "[Pt]": 627,
+    "[Ra]": 187,
+    "[Rb+1]": 207,
+    "[Rb]": 321,
+    "[Rh+2]": 628,
+    "[Ring1]": 9,
+    "[Ring1][#Branch1]": 344,
+    "[Ring1][#Branch2]": 353,
+    "[Ring1][#C]": 369,
+    "[Ring1][=Branch1]": 338,
+    "[Ring1][=Branch2]": 356,
+    "[Ring1][=C]": 363,
+    "[Ring1][=N]": 367,
+    "[Ring1][=O]": 550,
+    "[Ring1][Branch1]": 355,
+    "[Ring1][Branch2]": 357,
+    "[Ring1][C]": 373,
+    "[Ring1][N]": 368,
+    "[Ring1][O]": 354,
+    "[Ring1][P]": 372,
+    "[Ring1][Ring1]": 421,
+    "[Ring1][Ring2]": 439,
+    "[Ring1][S]": 365,
+    "[Ring2]": 12,
+    "[Ring2][=C]": 381,
+    "[Ring2][=N]": 459,
+    "[Ring2][=O]": 525,
+    "[Ring2][Branch1]": 520,
+    "[Ring2][C]": 393,
+    "[Ring2][N]": 433,
+    "[Ring2][O]": 466,
+    "[Ring2][Ring1]": 337,
+    "[Ring2][Ring2]": 359,
+    "[Ring2][S]": 502,
+    "[Ring3]": 629,
+    "[S+1]": 44,
+    "[S-1]": 60,
+    "[S-2]": 172,
+    "[S@+1]": 56,
+    "[S@@+1]": 95,
+    "[S@@]": 51,
+    "[S@]": 184,
+    "[SH0]": 170,
+    "[SH1+1]": 263,
+    "[SH1-1]": 205,
+    "[SH1]": 264,
+    "[SH2]": 299,
+    "[S]": 17,
+    "[S][=C]": 490,
+    "[S][=N]": 523,
+    "[S][=O]": 554,
+    "[S][Branch1]": 444,
+    "[S][C]": 375,
+    "[S][N]": 452,
+    "[S][O]": 527,
+    "[S][Ring1]": 485,
+    "[S][S]": 463,
+    "[Sb]": 211,
+    "[Se+1]": 230,
+    "[SeH1]": 130,
+    "[SeH2]": 330,
+    "[Se]": 36,
+    "[Si-1]": 331,
+    "[Si@]": 287,
+    "[SiH1-1]": 192,
+    "[SiH1]": 104,
+    "[SiH2]": 300,
+    "[SiH3-1]": 219,
+    "[SiH4]": 271,
+    "[Si]": 49,
+    "[Sn]": 630,
+    "[Sr+2]": 158,
+    "[Sr]": 332,
+    "[TeH1]": 279,
+    "[TeH2]": 239,
+    "[Te]": 81,
+    "[Xe]": 320,
+    "[Yb]": 209,
+    "[Zn+1]": 213,
+    "[Zn+2]": 148,
+    "[Zn-2]": 261,
+    "[Zn]": 82,
+    "[\\11CH3]": 203,
+    "[\\123I]": 214,
+    "[\\125I]": 79,
+    "[\\3H]": 250,
+    "[\\B]": 217,
+    "[\\Br]": 89,
+    "[\\C-1]": 265,
+    "[\\C@@H1]": 46,
+    "[\\C@@]": 80,
+    "[\\C@H1]": 77,
+    "[\\C@]": 99,
+    "[\\CH1-1]": 241,
+    "[\\C]": 43,
+    "[\\C][=C]": 436,
+    "[\\C][C]": 567,
+    "[\\Cl]": 58,
+    "[\\F]": 118,
+    "[\\I]": 59,
+    "[\\N+1]": 113,
+    "[\\N-1]": 222,
+    "[\\NH1]": 108,
+    "[\\N]": 55,
+    "[\\N][C]": 541,
+    "[\\O-1]": 64,
+    "[\\O]": 47,
+    "[\\P]": 93,
+    "[\\S+1]": 121,
+    "[\\S-1]": 255,
+    "[\\S]": 54,
+    "[\\Se]": 168,
+    "[\\Si]": 195
+}

tokenization_ape.py

@@ -17,7 +17,11 @@ from transformers import PreTrainedTokenizer
 
 Representation = Literal["SELFIES", "SMILES"]
 
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "selfies_vocab_file": "selfies_vocab.json",
+    "smiles_vocab_file": "smiles_vocab.json",
+}
 SELFIES_RE = re.compile(r"\[[^\]]+\]")
 SMILES_RE = re.compile(
     r"(\[[^\]]+\]|Br?|Cl?|Si?|Se?|Li?|Na?|Mg?|Al?|Ca?|Fe?|Zn?|"
@@ -64,6 +68,20 @@ def _normalize_representation(representation: str) -> Representation:
     return normalized  # type: ignore[return-value]
 
 
+def _select_vocab_file(
+    *,
+    representation: Representation,
+    vocab_file: str | os.PathLike[str] | None,
+    selfies_vocab_file: str | os.PathLike[str] | None,
+    smiles_vocab_file: str | os.PathLike[str] | None,
+) -> str | os.PathLike[str] | None:
+    if representation == "SELFIES" and selfies_vocab_file is not None:
+        return selfies_vocab_file
+    if representation == "SMILES" and smiles_vocab_file is not None:
+        return smiles_vocab_file
+    return vocab_file
+
+
 def pre_tokenize_molecule(molecule: str, representation: str) -> list[str]:
     active_representation = _normalize_representation(representation)
     if active_representation == "SELFIES":
@@ -97,19 +115,22 @@ def ape_tokenize(
 
     n = len(pieces)
     tokens: list[str] = []
+    append_token = tokens.append
+    vocab_contains = vocab.__contains__
+    join_pieces = "".join
     i = 0
 
     while i < n:
         upper = min(n, i + max_piece_span)
 
         for j in range(upper, i, -1):
-            candidate = "".join(pieces[i:j])
-            if candidate in vocab:
-                tokens.append(candidate)
+            candidate = join_pieces(pieces[i:j])
+            if vocab_contains(candidate):
+                append_token(candidate)
                 i = j
                 break
         else:
-            tokens.append(unk_token)
+            append_token(unk_token)
             i += 1
 
     return tokens
@@ -124,6 +145,8 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
     def __init__(
         self,
         vocab_file: str | os.PathLike[str] | None = None,
+        selfies_vocab_file: str | os.PathLike[str] | None = None,
+        smiles_vocab_file: str | os.PathLike[str] | None = None,
         vocab: dict[str, Any] | None = None,
         representation: str = "SELFIES",
         bos_token: str = "<s>",
@@ -134,8 +157,16 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
         model_max_length: int = 256,
         **kwargs,
     ) -> None:
+        self.representation = _normalize_representation(representation)
+        active_vocab_file = _select_vocab_file(
+            representation=self.representation,
+            vocab_file=vocab_file,
+            selfies_vocab_file=selfies_vocab_file,
+            smiles_vocab_file=smiles_vocab_file,
+        )
+
         if vocab is None:
-            if vocab_file is None:
+            if active_vocab_file is None:
                 vocab = {
                     bos_token: 0,
                     pad_token: 1,
@@ -144,13 +175,17 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
                     mask_token: 4,
                 }
             else:
-                with open(vocab_file, encoding="utf-8") as f:
+                with open(active_vocab_file, encoding="utf-8") as f:
                     vocab = json.load(f)
 
         if vocab is None:
             raise ValueError("Loaded vocabulary is None.")
 
-        self.vocab_file = str(vocab_file) if vocab_file is not None else None
+        self.vocab_file = str(active_vocab_file) if active_vocab_file is not None else None
+        self.selfies_vocab_file = (
+            str(selfies_vocab_file) if selfies_vocab_file is not None else None
+        )
+        self.smiles_vocab_file = str(smiles_vocab_file) if smiles_vocab_file is not None else None
         self.vocab = _coerce_vocab(vocab)
         self._require_special_tokens(
             bos_token=bos_token,
@@ -160,9 +195,8 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
             mask_token=mask_token,
         )
         self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
-        self.representation = _normalize_representation(representation)
         self.vocabulary_frequency: dict[str, int] = {}
-        self.pair_counts: dict[tuple[str, str] | str, int] = {}
+        self.pair_counts: dict[tuple[str, str], int] = {}
         self._max_piece_span = _max_vocab_piece_span(self.vocab, self.representation)
 
         super().__init__(
@@ -497,27 +531,49 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
         checkpoint_path: str = "checkpoint",
         checkpoint_interval: int = 500,
     ) -> None:
-        self.representation = _normalize_representation(representation or type)
-        text_padding = " " * 80
+        import warnings
+
+        new_rep = _normalize_representation(representation or type)
+        if new_rep != self.representation:
+            warnings.warn(
+                f"train() representation={new_rep!r} differs from tokenizer "
+                f"representation={self.representation!r}. Overwriting.",
+                UserWarning,
+                stacklevel=2,
+            )
+        self.representation = new_rep
 
-        print(f"Pretokenizing {self.representation}", end="\r")
-        tokenized_corpus = [
-            tokens for sentence in corpus if (tokens := self.pre_tokenize(str(sentence)))
-        ]
+        if not corpus:
+            raise ValueError("Cannot train APE tokenizer on an empty corpus.")
+
+        print(f"Pretokenizing {self.representation}...", flush=True)
+        tokenized_corpus = []
         vocabulary_frequency: defaultdict[str, int] = defaultdict(int)
-        for tokens in tokenized_corpus:
+        saw_tokens = False
+
+        for sentence in corpus:
+            tokens = self.pre_tokenize(str(sentence))
+            if not tokens:
+                continue
+            saw_tokens = True
             for token in tokens:
                 vocabulary_frequency[token] += 1
+            if len(tokens) > 1:
+                tokenized_corpus.append(tokens)
         print(
             f"Pretokenization complete, found {len(vocabulary_frequency)} tokens",
-            end="\r",
+            flush=True,
         )
 
-        if not tokenized_corpus:
+        if not saw_tokens:
             raise ValueError("Cannot train APE tokenizer on an empty corpus.")
 
         pre_tokens_counts = len(vocabulary_frequency)
         merged_counter = len(vocabulary_frequency) + 1
+        if save_checkpoint and checkpoint_interval <= 0:
+            raise ValueError(
+                "checkpoint_interval must be positive when save_checkpoint is enabled."
+            )
         checkpoint_increment = checkpoint_interval
         batch = checkpoint_interval + pre_tokens_counts
         piece_count_cache: dict[str, int] = {}
@@ -536,22 +592,26 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
                     pair = (tokens[i], tokens[i + 1])
 
                     if max_merge_pieces is not None:
-                        merged_candidate = "".join(pair)
+                        merged_candidate = pair[0] + pair[1]
                         if merged_piece_count(merged_candidate) > max_merge_pieces:
                             continue
 
                     pair_counts[pair] += 1
 
-            merged_pair_counts: dict[tuple[str, str] | str, int] = {
-                pair: count for pair, count in pair_counts.items()
-            }
-            self.pair_counts = merged_pair_counts
+            self.pair_counts = dict(pair_counts)
             if not pair_counts:
                 return ("", ""), 0
-            return max(pair_counts.items(), key=lambda x: x[1], default=(("", ""), 0))
+
+            most_common_pair = ("", "")
+            most_common_frequency = 0
+            for pair, count in pair_counts.items():
+                if count > most_common_frequency:
+                    most_common_pair = pair
+                    most_common_frequency = count
+            return most_common_pair, most_common_frequency
 
         while True:
-            if save_checkpoint and len(vocabulary_frequency) == batch:
+            if save_checkpoint and len(vocabulary_frequency) >= batch:
                 self.vocabulary_frequency = dict(vocabulary_frequency)
                 self.vocab = {
                     **{
@@ -579,51 +639,53 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
                 batch += checkpoint_increment
 
             if len(vocabulary_frequency) >= max_vocab_size:
-                print("\rMax vocabulary achieved", text_padding)
+                print("Max vocabulary achieved", flush=True)
                 break
 
-            if all(len(tokens) < 2 for tokens in tokenized_corpus):
-                print("\rNo more mergeable pairs", text_padding)
+            if not tokenized_corpus:
+                print("No more mergeable pairs", flush=True)
                 break
 
             most_common_pair, freq = get_most_common_pair(tokenized_corpus)
             if freq < min_freq_for_merge:
-                print("\rNot enough frequency found", text_padding)
+                print("Not enough frequency found", flush=True)
                 break
 
             if not most_common_pair[0] or not most_common_pair[1]:
-                print("\rNo valid merge pair found", text_padding)
+                print("No valid merge pair found", flush=True)
                 break
 
-            merged_word = "".join(most_common_pair)
+            left_token, right_token = most_common_pair
+            merged_word = left_token + right_token
             if merged_word not in vocabulary_frequency:
                 print(
                     f"New merge found: {merged_word} {merged_counter}/{max_vocab_size} "
-                    f"{round(merged_counter / max_vocab_size * 100, 2)}%"
+                    f"{round(merged_counter / max_vocab_size * 100, 2)}%",
+                    flush=True,
                 )
                 merged_counter += 1
-            vocabulary_frequency[merged_word] = vocabulary_frequency.get(merged_word, 0) + freq
+            vocabulary_frequency[merged_word] += freq
 
             new_tokenized_corpus = []
             for tokens in tokenized_corpus:
                 new_tokens = []
-                skip_next = False
-                for i in range(len(tokens)):
-                    if skip_next:
-                        skip_next = False
-                        continue
-
+                append_token = new_tokens.append
+                i = 0
+                token_count = len(tokens)
+                while i < token_count:
                     if (
-                        i < len(tokens) - 1
-                        and tokens[i] == most_common_pair[0]
-                        and tokens[i + 1] == most_common_pair[1]
+                        i < token_count - 1
+                        and tokens[i] == left_token
+                        and tokens[i + 1] == right_token
                     ):
-                        new_tokens.append(merged_word)
-                        skip_next = True
+                        append_token(merged_word)
+                        i += 2
                     else:
-                        new_tokens.append(tokens[i])
+                        append_token(tokens[i])
+                        i += 1
 
-                new_tokenized_corpus.append(new_tokens)
+                if len(new_tokens) > 1:
+                    new_tokenized_corpus.append(new_tokens)
 
             tokenized_corpus = new_tokenized_corpus
 
@@ -640,8 +702,6 @@ class APEPreTrainedTokenizer(PreTrainedTokenizer):
         self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
         self._refresh_tokenization_cache()
 
-        checkpoint_dir = Path(checkpoint_path)
-
     def train_from_iterator(self, iterator, *args, **kwargs) -> None:
         raise NotImplementedError("train_from_iterator is not implemented for APE")
 

tokenizer_config.json

@@ -51,8 +51,13 @@
   "bos_token": "<s>",
   "eos_token": "</s>",
   "mask_token": "<mask>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
   "model_max_length": 256,
   "pad_token": "<pad>",
   "representation": "SELFIES",
-  "unk_token": "<unk>"
-}
+  "unk_token": "<unk>",
+  "use_fast": false
+}

tokenizer_metadata.json

@@ -0,0 +1,28 @@
+{
+  "ape_source": "modernmolbert.local",
+  "created_at_utc": "2026-05-19T16:24:35.971419+00:00",
+  "creation_command": "python -m modernmolbert.train_ape_tokenizer",
+  "dataset_name": "data/pretrain/chembl36_selfies",
+  "extra_vocab_selfies_path": null,
+  "extra_vocab_symbols_added": 42,
+  "extra_vocab_symbols_path": "tokenizer/extra_symbols/benchmark_missing_selfies_symbols_min10.txt",
+  "extra_vocab_symbols_requested": 42,
+  "max_merge_pieces": 2,
+  "max_vocab_size": 2000,
+  "min_freq_for_merge": 3000,
+  "representation": "SELFIES",
+  "seed": 42,
+  "selfies_column": "selfies",
+  "shuffle_buffer_size": 100000,
+  "special_ids": {
+    "bos_token": 0,
+    "eos_token": 2,
+    "mask_token": 4,
+    "pad_token": 1,
+    "unk_token": 3
+  },
+  "tokenizer_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json",
+  "tokenizer_sha256": "26ad0e90de9c0a469eb6e3b7aa985e2a2d7cc3f0dd7b2b229ccfce41f639e208",
+  "tokenizer_train_size": 2000000,
+  "vocab_size": 631
+}

train_results.json

@@ -0,0 +1,10 @@
+{
+    "epoch": 3.066,
+    "num_parameters": 34149495.0,
+    "total_flos": 9.74439612930639e+16,
+    "train_loss": 0.6228721523284912,
+    "train_runtime": 2832.8619,
+    "train_samples_per_second": 2711.039,
+    "train_samples_streaming": 7680000.0,
+    "train_steps_per_second": 10.59
+}

trainer_state.json

@@ -1,7 +1,7 @@
 {
-  "best_global_step": 25000,
-  "best_metric": 1.500394344329834,
-  "best_model_checkpoint": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_1e-4/checkpoint-25000",
+  "best_global_step": 30000,
+  "best_metric": 0.37694016098976135,
+  "best_model_checkpoint": "runs/chembl36_small_mask_mlm_lr_sweep/mask_standard__mlm_0p15__lr_4e-4/checkpoint-30000",
   "epoch": 3.066,
   "eval_steps": 5000,
   "global_step": 30000,
@@ -11,2174 +11,2174 @@
   "log_history": [
     {
       "epoch": 0.0033333333333333335,
-      "grad_norm": 4.756051063537598,
-      "learning_rate": 6.6e-06,
-      "loss": 7.7740185546875,
+      "grad_norm": 2.7403433322906494,
+      "learning_rate": 2.64e-05,
+      "loss": 5.425458984375,
       "step": 100
     },
     {
       "epoch": 0.006666666666666667,
-      "grad_norm": 4.41278076171875,
-      "learning_rate": 1.3266666666666666e-05,
-      "loss": 7.52440673828125,
+      "grad_norm": 4.65195894241333,
+      "learning_rate": 5.3066666666666665e-05,
+      "loss": 3.9814361572265624,
       "step": 200
     },
     {
       "epoch": 0.01,
-      "grad_norm": 3.1706013679504395,
-      "learning_rate": 1.9933333333333334e-05,
-      "loss": 7.2918212890625,
+      "grad_norm": 2.9648351669311523,
+      "learning_rate": 7.973333333333334e-05,
+      "loss": 2.9717633056640627,
       "step": 300
     },
     {
       "epoch": 0.013333333333333334,
-      "grad_norm": 3.3350961208343506,
-      "learning_rate": 2.6600000000000003e-05,
-      "loss": 7.060899047851563,
+      "grad_norm": 2.79247784614563,
+      "learning_rate": 0.00010640000000000001,
+      "loss": 2.318418426513672,
       "step": 400
     },
     {
       "epoch": 0.016666666666666666,
-      "grad_norm": 3.3246028423309326,
-      "learning_rate": 3.326666666666667e-05,
-      "loss": 6.882850341796875,
+      "grad_norm": 2.95702862739563,
+      "learning_rate": 0.00013306666666666668,
+      "loss": 2.012203826904297,
       "step": 500
     },
     {
       "epoch": 0.02,
-      "grad_norm": 3.528203248977661,
-      "learning_rate": 3.993333333333333e-05,
-      "loss": 6.68141357421875,
+      "grad_norm": 2.464130401611328,
+      "learning_rate": 0.00015973333333333333,
+      "loss": 1.8253421020507812,
       "step": 600
     },
     {
       "epoch": 0.023333333333333334,
-      "grad_norm": 3.592036724090576,
-      "learning_rate": 4.660000000000001e-05,
-      "loss": 6.46164794921875,
+      "grad_norm": 2.5683345794677734,
+      "learning_rate": 0.00018640000000000003,
+      "loss": 1.7040005493164063,
       "step": 700
     },
     {
       "epoch": 0.02666666666666667,
-      "grad_norm": 3.5415823459625244,
-      "learning_rate": 5.326666666666666e-05,
-      "loss": 6.2806982421875,
+      "grad_norm": 2.7064743041992188,
+      "learning_rate": 0.00021306666666666665,
+      "loss": 1.5912832641601562,
       "step": 800
     },
     {
       "epoch": 0.03,
-      "grad_norm": 3.8730804920196533,
-      "learning_rate": 5.9933333333333345e-05,
-      "loss": 6.08379638671875,
+      "grad_norm": 1.7914412021636963,
+      "learning_rate": 0.00023973333333333338,
+      "loss": 1.532608184814453,
       "step": 900
     },
     {
       "epoch": 0.03333333333333333,
-      "grad_norm": 4.294698715209961,
-      "learning_rate": 6.66e-05,
-      "loss": 5.890241088867188,
+      "grad_norm": 2.02378511428833,
+      "learning_rate": 0.0002664,
+      "loss": 1.4479930114746093,
       "step": 1000
     },
     {
       "epoch": 0.03666666666666667,
-      "grad_norm": 4.327278137207031,
-      "learning_rate": 7.326666666666667e-05,
-      "loss": 5.653626098632812,
+      "grad_norm": 1.8761825561523438,
+      "learning_rate": 0.00029306666666666667,
+      "loss": 1.4139170837402344,
       "step": 1100
     },
     {
       "epoch": 0.04,
-      "grad_norm": 4.361931800842285,
-      "learning_rate": 7.993333333333334e-05,
-      "loss": 5.426879272460938,
+      "grad_norm": 1.751976490020752,
+      "learning_rate": 0.00031973333333333337,
+      "loss": 1.3751948547363282,
       "step": 1200
     },
     {
       "epoch": 0.043333333333333335,
-      "grad_norm": 4.677211284637451,
-      "learning_rate": 8.66e-05,
-      "loss": 5.20781982421875,
+      "grad_norm": 1.8251773118972778,
+      "learning_rate": 0.0003464,
+      "loss": 1.3400251770019531,
       "step": 1300
     },
     {
       "epoch": 0.04666666666666667,
-      "grad_norm": 4.3806867599487305,
-      "learning_rate": 9.326666666666667e-05,
-      "loss": 4.999788818359375,
+      "grad_norm": 1.5843342542648315,
+      "learning_rate": 0.00037306666666666666,
+      "loss": 1.327974853515625,
       "step": 1400
     },
     {
       "epoch": 0.05,
-      "grad_norm": 4.898468494415283,
-      "learning_rate": 9.993333333333334e-05,
-      "loss": 4.8133847045898435,
+      "grad_norm": 1.5720994472503662,
+      "learning_rate": 0.00039973333333333336,
+      "loss": 1.2728753662109376,
       "step": 1500
     },
     {
       "epoch": 0.05333333333333334,
-      "grad_norm": 4.611327648162842,
-      "learning_rate": 9.999702274445236e-05,
-      "loss": 4.6433349609375,
+      "grad_norm": 1.2932488918304443,
+      "learning_rate": 0.0003999880909778094,
+      "loss": 1.259843521118164,
       "step": 1600
     },
     {
       "epoch": 0.056666666666666664,
-      "grad_norm": 5.104644298553467,
-      "learning_rate": 9.998797074387361e-05,
-      "loss": 4.487667846679687,
+      "grad_norm": 1.4394837617874146,
+      "learning_rate": 0.00039995188297549443,
+      "loss": 1.2250770568847655,
       "step": 1700
     },
     {
       "epoch": 0.06,
-      "grad_norm": 4.787322521209717,
-      "learning_rate": 9.997284479438253e-05,
-      "loss": 4.340900573730469,
+      "grad_norm": 1.3632104396820068,
+      "learning_rate": 0.0003998913791775301,
+      "loss": 1.1869202423095704,
       "step": 1800
     },
     {
       "epoch": 0.06333333333333334,
-      "grad_norm": 4.701177597045898,
-      "learning_rate": 9.995164673390625e-05,
-      "loss": 4.235370178222656,
+      "grad_norm": 1.4223003387451172,
+      "learning_rate": 0.000399806586935625,
+      "loss": 1.1602543640136718,
       "step": 1900
     },
     {
       "epoch": 0.06666666666666667,
-      "grad_norm": 4.965324878692627,
-      "learning_rate": 9.992437913818312e-05,
-      "loss": 4.140614013671875,
+      "grad_norm": 1.3835265636444092,
+      "learning_rate": 0.00039969751655273246,
+      "loss": 1.1331352233886718,
       "step": 2000
     },
     {
       "epoch": 0.07,
-      "grad_norm": 5.183881759643555,
-      "learning_rate": 9.989104532044994e-05,
-      "loss": 4.013364868164063,
+      "grad_norm": 1.4172451496124268,
+      "learning_rate": 0.0003995641812817998,
+      "loss": 1.1115520477294922,
       "step": 2100
     },
     {
       "epoch": 0.07333333333333333,
-      "grad_norm": 5.037045001983643,
-      "learning_rate": 9.985164933103929e-05,
-      "loss": 3.9139785766601562,
+      "grad_norm": 1.2611116170883179,
+      "learning_rate": 0.00039940659732415714,
+      "loss": 1.0953889465332032,
       "step": 2200
     },
     {
       "epoch": 0.07666666666666666,
-      "grad_norm": 5.492240905761719,
-      "learning_rate": 9.980619595688737e-05,
-      "loss": 3.8507009887695314,
+      "grad_norm": 1.4010944366455078,
+      "learning_rate": 0.00039922478382754947,
+      "loss": 1.0633756256103515,
       "step": 2300
     },
     {
       "epoch": 0.08,
-      "grad_norm": 5.084971904754639,
-      "learning_rate": 9.975469072095237e-05,
-      "loss": 3.756464538574219,
+      "grad_norm": 1.277366280555725,
+      "learning_rate": 0.00039901876288380946,
+      "loss": 1.0482559967041016,
       "step": 2400
     },
     {
       "epoch": 0.08333333333333333,
-      "grad_norm": 5.1968793869018555,
-      "learning_rate": 9.969713988154339e-05,
-      "loss": 3.6725424194335936,
+      "grad_norm": 1.1845537424087524,
+      "learning_rate": 0.00039878855952617357,
+      "loss": 1.0299220275878906,
       "step": 2500
     },
     {
       "epoch": 0.08666666666666667,
-      "grad_norm": 5.174182891845703,
-      "learning_rate": 9.963355043155996e-05,
-      "loss": 3.606795349121094,
+      "grad_norm": 1.1772642135620117,
+      "learning_rate": 0.00039853420172623985,
+      "loss": 1.0119639587402345,
       "step": 2600
     },
     {
       "epoch": 0.09,
-      "grad_norm": 5.094499588012695,
-      "learning_rate": 9.956393009764244e-05,
-      "loss": 3.534298095703125,
+      "grad_norm": 1.241080641746521,
+      "learning_rate": 0.00039825572039056976,
+      "loss": 0.9945954895019531,
       "step": 2700
     },
     {
       "epoch": 0.09333333333333334,
-      "grad_norm": 5.067377090454102,
-      "learning_rate": 9.948828733923305e-05,
-      "loss": 3.488475036621094,
+      "grad_norm": 1.1556016206741333,
+      "learning_rate": 0.0003979531493569322,
+      "loss": 0.9795339965820312,
       "step": 2800
     },
     {
       "epoch": 0.09666666666666666,
-      "grad_norm": 5.349589824676514,
-      "learning_rate": 9.94066313475481e-05,
-      "loss": 3.4123611450195312,
+      "grad_norm": 1.1028602123260498,
+      "learning_rate": 0.0003976265253901924,
+      "loss": 0.9684971618652344,
       "step": 2900
     },
     {
       "epoch": 0.1,
-      "grad_norm": 4.99954080581665,
-      "learning_rate": 9.931897204446104e-05,
-      "loss": 3.3717422485351562,
+      "grad_norm": 1.3073146343231201,
+      "learning_rate": 0.00039727588817784414,
+      "loss": 0.9519919586181641,
       "step": 3000
     },
     {
       "epoch": 0.10333333333333333,
-      "grad_norm": 5.293025970458984,
-      "learning_rate": 9.9225320081297e-05,
-      "loss": 3.2826019287109376,
+      "grad_norm": 1.224905252456665,
+      "learning_rate": 0.000396901280325188,
+      "loss": 0.9399765777587891,
       "step": 3100
     },
     {
       "epoch": 0.10666666666666667,
-      "grad_norm": 4.9904375076293945,
-      "learning_rate": 9.912568683753853e-05,
-      "loss": 3.269647521972656,
+      "grad_norm": 1.0792871713638306,
+      "learning_rate": 0.0003965027473501541,
+      "loss": 0.9271878814697265,
       "step": 3200
     },
     {
       "epoch": 0.11,
-      "grad_norm": 4.861743450164795,
-      "learning_rate": 9.902008441944286e-05,
-      "loss": 3.2435989379882812,
+      "grad_norm": 1.077876329421997,
+      "learning_rate": 0.0003960803376777714,
+      "loss": 0.9128645324707031,
       "step": 3300
     },
     {
       "epoch": 0.11333333333333333,
-      "grad_norm": 5.036742210388184,
-      "learning_rate": 9.890852565857092e-05,
-      "loss": 3.1625836181640623,
+      "grad_norm": 1.1184345483779907,
+      "learning_rate": 0.0003956341026342837,
+      "loss": 0.9120386505126953,
       "step": 3400
     },
     {
       "epoch": 0.11666666666666667,
-      "grad_norm": 5.042503356933594,
-      "learning_rate": 9.879102411022817e-05,
-      "loss": 3.1680789184570313,
+      "grad_norm": 1.1383893489837646,
+      "learning_rate": 0.0003951640964409127,
+      "loss": 0.9151699829101563,
       "step": 3500
     },
     {
       "epoch": 0.12,
-      "grad_norm": 5.01936149597168,
-      "learning_rate": 9.866759405181765e-05,
-      "loss": 3.1387368774414064,
+      "grad_norm": 1.12523353099823,
+      "learning_rate": 0.0003946703762072706,
+      "loss": 0.8922748565673828,
       "step": 3600
     },
     {
       "epoch": 0.12333333333333334,
-      "grad_norm": 5.504270553588867,
-      "learning_rate": 9.85382504811049e-05,
-      "loss": 3.1010336303710937,
+      "grad_norm": 1.0355699062347412,
+      "learning_rate": 0.0003941530019244196,
+      "loss": 0.8890122222900391,
       "step": 3700
     },
     {
       "epoch": 0.12666666666666668,
-      "grad_norm": 5.191413402557373,
-      "learning_rate": 9.840300911439591e-05,
-      "loss": 3.057752990722656,
+      "grad_norm": 1.0852240324020386,
+      "learning_rate": 0.00039361203645758364,
+      "loss": 0.8781277465820313,
       "step": 3800
     },
     {
       "epoch": 0.13,
-      "grad_norm": 5.080201625823975,
-      "learning_rate": 9.826188638462718e-05,
-      "loss": 3.043381652832031,
+      "grad_norm": 1.0739648342132568,
+      "learning_rate": 0.00039304754553850873,
+      "loss": 0.8778704833984375,
       "step": 3900
     },
     {
       "epoch": 0.13333333333333333,
-      "grad_norm": 5.173981189727783,
-      "learning_rate": 9.811489943936922e-05,
-      "loss": 3.011300354003906,
+      "grad_norm": 1.0522024631500244,
+      "learning_rate": 0.0003924595977574769,
+      "loss": 0.8774790954589844,
       "step": 4000
     },
     {
       "epoch": 0.13666666666666666,
-      "grad_norm": 5.164648532867432,
-      "learning_rate": 9.796206613874283e-05,
-      "loss": 2.944804382324219,
+      "grad_norm": 1.0717562437057495,
+      "learning_rate": 0.00039184826455497133,
+      "loss": 0.8516897583007812,
       "step": 4100
     },
     {
       "epoch": 0.14,
-      "grad_norm": 5.0373969078063965,
-      "learning_rate": 9.780340505324901e-05,
-      "loss": 2.9307366943359376,
+      "grad_norm": 1.0378234386444092,
+      "learning_rate": 0.00039121362021299604,
+      "loss": 0.8448622894287109,
       "step": 4200
     },
     {
       "epoch": 0.14333333333333334,
-      "grad_norm": 4.772580623626709,
-      "learning_rate": 9.763893546151244e-05,
-      "loss": 2.9096554565429686,
+      "grad_norm": 1.0532786846160889,
+      "learning_rate": 0.00039055574184604976,
+      "loss": 0.8344584655761719,
       "step": 4300
     },
     {
       "epoch": 0.14666666666666667,
-      "grad_norm": 4.913202285766602,
-      "learning_rate": 9.746867734793903e-05,
-      "loss": 2.8842669677734376,
+      "grad_norm": 1.021437644958496,
+      "learning_rate": 0.0003898747093917561,
+      "loss": 0.8287362670898437,
       "step": 4400
     },
     {
       "epoch": 0.15,
-      "grad_norm": 5.291968822479248,
-      "learning_rate": 9.729265140028762e-05,
-      "loss": 2.8610394287109373,
+      "grad_norm": 0.9874396920204163,
+      "learning_rate": 0.0003891706056011505,
+      "loss": 0.8258264923095703,
       "step": 4500
     },
     {
       "epoch": 0.15333333333333332,
-      "grad_norm": 4.679079055786133,
-      "learning_rate": 9.711087900715627e-05,
-      "loss": 2.8442318725585936,
+      "grad_norm": 0.9322711825370789,
+      "learning_rate": 0.0003884435160286251,
+      "loss": 0.816080322265625,
       "step": 4600
     },
     {
       "epoch": 0.15666666666666668,
-      "grad_norm": 5.12017822265625,
-      "learning_rate": 9.692338225538333e-05,
-      "loss": 2.7991876220703125,
+      "grad_norm": 0.9402130842208862,
+      "learning_rate": 0.00038769352902153333,
+      "loss": 0.8135909271240235,
       "step": 4700
     },
     {
       "epoch": 0.16,
-      "grad_norm": 4.660444259643555,
-      "learning_rate": 9.673018392736374e-05,
-      "loss": 2.7571255493164064,
+      "grad_norm": 1.0193718671798706,
+      "learning_rate": 0.00038692073570945496,
+      "loss": 0.817694091796875,
       "step": 4800
     },
     {
       "epoch": 0.16333333333333333,
-      "grad_norm": 5.051372528076172,
-      "learning_rate": 9.653130749828075e-05,
-      "loss": 2.7532052612304687,
+      "grad_norm": 0.9926565289497375,
+      "learning_rate": 0.000386125229993123,
+      "loss": 0.7987307739257813,
       "step": 4900
     },
     {
       "epoch": 0.16666666666666666,
-      "grad_norm": 4.846753120422363,
-      "learning_rate": 9.632677713325353e-05,
-      "loss": 2.724853515625,
+      "grad_norm": 0.9364326596260071,
+      "learning_rate": 0.0003853071085330141,
+      "loss": 0.8007232666015625,
       "step": 5000
     },
     {
       "epoch": 0.16666666666666666,
-      "eval_loss": 2.7221035957336426,
-      "eval_masked_accuracy": 0.40898550724637683,
-      "eval_runtime": 5.8666,
-      "eval_samples_per_second": 698.194,
-      "eval_steps_per_second": 2.727,
+      "eval_loss": 0.7967580556869507,
+      "eval_masked_accuracy": 0.7653559496273451,
+      "eval_runtime": 2.5634,
+      "eval_samples_per_second": 1597.853,
+      "eval_steps_per_second": 6.242,
       "step": 5000
     },
     {
       "epoch": 0.17,
-      "grad_norm": 4.767246246337891,
-      "learning_rate": 9.611661768440093e-05,
-      "loss": 2.6953024291992187,
+      "grad_norm": 0.9981640577316284,
+      "learning_rate": 0.0003844664707376037,
+      "loss": 0.7831909942626953,
       "step": 5100
     },
     {
       "epoch": 0.17333333333333334,
-      "grad_norm": 4.842975616455078,
-      "learning_rate": 9.590085468782162e-05,
-      "loss": 2.6890545654296876,
+      "grad_norm": 0.935793936252594,
+      "learning_rate": 0.00038360341875128646,
+      "loss": 0.7719477081298828,
       "step": 5200
     },
     {
       "epoch": 0.17666666666666667,
-      "grad_norm": 4.971850872039795,
-      "learning_rate": 9.56795143604914e-05,
-      "loss": 2.634810791015625,
+      "grad_norm": 0.9942166805267334,
+      "learning_rate": 0.0003827180574419656,
+      "loss": 0.7788728332519531,
       "step": 5300
     },
     {
       "epoch": 0.18,
-      "grad_norm": 4.957698822021484,
-      "learning_rate": 9.545262359707756e-05,
-      "loss": 2.6262237548828127,
+      "grad_norm": 0.921302318572998,
+      "learning_rate": 0.00038181049438831024,
+      "loss": 0.7685958099365234,
       "step": 5400
     },
     {
       "epoch": 0.18333333333333332,
-      "grad_norm": 4.960840225219727,
-      "learning_rate": 9.522020996667092e-05,
-      "loss": 2.5996078491210937,
+      "grad_norm": 0.8787013292312622,
+      "learning_rate": 0.0003808808398666837,
+      "loss": 0.7563288116455078,
       "step": 5500
     },
     {
       "epoch": 0.18666666666666668,
-      "grad_norm": 4.946103096008301,
-      "learning_rate": 9.498230170943596e-05,
-      "loss": 2.6009527587890626,
+      "grad_norm": 0.8968636989593506,
+      "learning_rate": 0.00037992920683774386,
+      "loss": 0.7474672698974609,
       "step": 5600
     },
     {
       "epoch": 0.19,
-      "grad_norm": 4.791319847106934,
-      "learning_rate": 9.473892773317952e-05,
-      "loss": 2.5778823852539063,
+      "grad_norm": 0.9104163646697998,
+      "learning_rate": 0.00037895571093271807,
+      "loss": 0.7457213592529297,
       "step": 5700
     },
     {
       "epoch": 0.19333333333333333,
-      "grad_norm": 4.821264743804932,
-      "learning_rate": 9.449011760983809e-05,
-      "loss": 2.5622207641601564,
+      "grad_norm": 0.8929153680801392,
+      "learning_rate": 0.00037796047043935234,
+      "loss": 0.7423068237304687,
       "step": 5800
     },
     {
       "epoch": 0.19666666666666666,
-      "grad_norm": 4.989269256591797,
-      "learning_rate": 9.423590157188474e-05,
-      "loss": 2.527291564941406,
+      "grad_norm": 0.9810736775398254,
+      "learning_rate": 0.00037694360628753897,
+      "loss": 0.7346783447265625,
       "step": 5900
     },
     {
       "epoch": 0.2,
-      "grad_norm": 4.952500820159912,
-      "learning_rate": 9.397631050865554e-05,
-      "loss": 2.5371446228027343,
+      "grad_norm": 0.931078314781189,
+      "learning_rate": 0.00037590524203462216,
+      "loss": 0.7334535217285156,
       "step": 6000
     },
     {
       "epoch": 0.20333333333333334,
-      "grad_norm": 4.967984199523926,
-      "learning_rate": 9.371137596259623e-05,
-      "loss": 2.529341735839844,
+      "grad_norm": 1.019803524017334,
+      "learning_rate": 0.0003748455038503849,
+      "loss": 0.7276763916015625,
       "step": 6100
     },
     {
       "epoch": 0.20666666666666667,
-      "grad_norm": 4.9603352546691895,
-      "learning_rate": 9.344113012542963e-05,
-      "loss": 2.478318634033203,
+      "grad_norm": 0.9782306551933289,
+      "learning_rate": 0.0003737645205017185,
+      "loss": 0.7344593048095703,
       "step": 6200
     },
     {
       "epoch": 0.21,
-      "grad_norm": 4.8668389320373535,
-      "learning_rate": 9.316560583424403e-05,
-      "loss": 2.4962451171875,
+      "grad_norm": 0.9293437004089355,
+      "learning_rate": 0.0003726624233369761,
+      "loss": 0.7246415710449219,
       "step": 6300
     },
     {
       "epoch": 0.21333333333333335,
-      "grad_norm": 4.65514612197876,
-      "learning_rate": 9.288483656750322e-05,
-      "loss": 2.465853424072266,
+      "grad_norm": 0.8795522451400757,
+      "learning_rate": 0.0003715393462700129,
+      "loss": 0.7216864776611328,
       "step": 6400
     },
     {
       "epoch": 0.21666666666666667,
-      "grad_norm": 4.887564182281494,
-      "learning_rate": 9.259885644097861e-05,
-      "loss": 2.454242858886719,
+      "grad_norm": 0.9632502198219299,
+      "learning_rate": 0.00037039542576391445,
+      "loss": 0.7085445404052735,
       "step": 6500
     },
     {
       "epoch": 0.22,
-      "grad_norm": 4.730443477630615,
-      "learning_rate": 9.230770020360382e-05,
-      "loss": 2.443849639892578,
+      "grad_norm": 0.8495221138000488,
+      "learning_rate": 0.0003692308008144153,
+      "loss": 0.7115458679199219,
       "step": 6600
     },
     {
       "epoch": 0.22333333333333333,
-      "grad_norm": 5.151344299316406,
-      "learning_rate": 9.201140323325248e-05,
-      "loss": 2.4249298095703127,
+      "grad_norm": 0.8541926741600037,
+      "learning_rate": 0.0003680456129330099,
+      "loss": 0.6990865325927734,
       "step": 6700
     },
     {
       "epoch": 0.22666666666666666,
-      "grad_norm": 5.001748561859131,
-      "learning_rate": 9.171000153243947e-05,
-      "loss": 2.3924867248535158,
+      "grad_norm": 0.8344811201095581,
+      "learning_rate": 0.0003668400061297579,
+      "loss": 0.6993592071533203,
       "step": 6800
     },
     {
       "epoch": 0.23,
-      "grad_norm": 4.810725688934326,
-      "learning_rate": 9.140353172394636e-05,
-      "loss": 2.4176744079589843,
+      "grad_norm": 0.8975337147712708,
+      "learning_rate": 0.00036561412689578544,
+      "loss": 0.6972901153564454,
       "step": 6900
     },
     {
       "epoch": 0.23333333333333334,
-      "grad_norm": 4.894969463348389,
-      "learning_rate": 9.109203104637137e-05,
-      "loss": 2.3665452575683594,
+      "grad_norm": 0.8364660143852234,
+      "learning_rate": 0.00036436812418548547,
+      "loss": 0.6967823791503907,
       "step": 7000
     },
     {
       "epoch": 0.23666666666666666,
-      "grad_norm": 4.827425479888916,
-      "learning_rate": 9.077553734960468e-05,
-      "loss": 2.3609941101074217,
+      "grad_norm": 0.8676795363426208,
+      "learning_rate": 0.00036310214939841873,
+      "loss": 0.6893701934814453,
       "step": 7100
     },
     {
       "epoch": 0.24,
-      "grad_norm": 4.975261211395264,
-      "learning_rate": 9.045408909022928e-05,
-      "loss": 2.3557579040527346,
+      "grad_norm": 0.8879637718200684,
+      "learning_rate": 0.0003618163563609171,
+      "loss": 0.6855775451660157,
       "step": 7200
     },
     {
       "epoch": 0.24333333333333335,
-      "grad_norm": 4.766048908233643,
-      "learning_rate": 9.012772532684818e-05,
-      "loss": 2.3431201171875,
+      "grad_norm": 0.9070082306861877,
+      "learning_rate": 0.00036051090130739274,
+      "loss": 0.6840663146972656,
       "step": 7300
     },
     {
       "epoch": 0.24666666666666667,
-      "grad_norm": 4.788959503173828,
-      "learning_rate": 8.979648571533852e-05,
-      "loss": 2.3297906494140626,
+      "grad_norm": 0.8394324779510498,
+      "learning_rate": 0.0003591859428613541,
+      "loss": 0.6798997497558594,
       "step": 7400
     },
     {
       "epoch": 0.25,
-      "grad_norm": 4.961067199707031,
-      "learning_rate": 8.9460410504033e-05,
-      "loss": 2.331878662109375,
+      "grad_norm": 0.8838658332824707,
+      "learning_rate": 0.000357841642016132,
+      "loss": 0.6784028625488281,
       "step": 7500
     },
     {
       "epoch": 0.25333333333333335,
-      "grad_norm": 4.572299003601074,
-      "learning_rate": 8.91195405288294e-05,
-      "loss": 2.29117919921875,
+      "grad_norm": 0.8179841041564941,
+      "learning_rate": 0.0003564781621153176,
+      "loss": 0.6701314544677734,
       "step": 7600
     },
     {
       "epoch": 0.25666666666666665,
-      "grad_norm": 4.729584217071533,
-      "learning_rate": 8.877391720822875e-05,
-      "loss": 2.2867362976074217,
+      "grad_norm": 0.8342909812927246,
+      "learning_rate": 0.000355095668832915,
+      "loss": 0.6615406799316407,
       "step": 7700
     },
     {
       "epoch": 0.26,
-      "grad_norm": 4.712250232696533,
-      "learning_rate": 8.842358253830246e-05,
-      "loss": 2.2930256652832033,
+      "grad_norm": 0.8877819180488586,
+      "learning_rate": 0.00035369433015320983,
+      "loss": 0.6670161437988281,
       "step": 7800
     },
     {
       "epoch": 0.2633333333333333,
-      "grad_norm": 4.344603061676025,
-      "learning_rate": 8.806857908758967e-05,
-      "loss": 2.2762042236328126,
+      "grad_norm": 0.9063718914985657,
+      "learning_rate": 0.0003522743163503587,
+      "loss": 0.6525628662109375,
       "step": 7900
     },
     {
       "epoch": 0.26666666666666666,
-      "grad_norm": 4.702202320098877,
-      "learning_rate": 8.770894999192467e-05,
-      "loss": 2.2673519897460936,
+      "grad_norm": 0.8738296627998352,
+      "learning_rate": 0.00035083579996769866,
+      "loss": 0.6613227081298828,
       "step": 8000
     },
     {
       "epoch": 0.27,
-      "grad_norm": 4.727926731109619,
-      "learning_rate": 8.734473894919564e-05,
-      "loss": 2.2311663818359375,
+      "grad_norm": 0.9108850359916687,
+      "learning_rate": 0.00034937895579678256,
+      "loss": 0.6542455291748047,
       "step": 8100
     },
     {
       "epoch": 0.2733333333333333,
-      "grad_norm": 4.6895294189453125,
-      "learning_rate": 8.697599021403489e-05,
-      "loss": 2.2383134460449217,
+      "grad_norm": 0.7989206314086914,
+      "learning_rate": 0.00034790396085613955,
+      "loss": 0.6540352630615235,
       "step": 8200
     },
     {
       "epoch": 0.27666666666666667,
-      "grad_norm": 4.4664998054504395,
-      "learning_rate": 8.660274859244168e-05,
-      "loss": 2.1981523132324217,
+      "grad_norm": 0.8691635727882385,
+      "learning_rate": 0.0003464109943697667,
+      "loss": 0.6410691070556641,
       "step": 8300
     },
     {
       "epoch": 0.28,
-      "grad_norm": 4.363753318786621,
-      "learning_rate": 8.62250594363378e-05,
-      "loss": 2.1935469055175782,
+      "grad_norm": 0.8667342066764832,
+      "learning_rate": 0.0003449002377453512,
+      "loss": 0.6395606994628906,
       "step": 8400
     },
     {
       "epoch": 0.2833333333333333,
-      "grad_norm": 4.9684062004089355,
-      "learning_rate": 8.584296863805709e-05,
-      "loss": 2.1947981262207032,
+      "grad_norm": 0.7986162900924683,
+      "learning_rate": 0.00034337187455222835,
+      "loss": 0.6376551055908203,
       "step": 8500
     },
     {
       "epoch": 0.2866666666666667,
-      "grad_norm": 4.717839241027832,
-      "learning_rate": 8.545652262476897e-05,
-      "loss": 2.190624237060547,
+      "grad_norm": 0.9031108617782593,
+      "learning_rate": 0.0003418260904990759,
+      "loss": 0.6303312683105469,
       "step": 8600
     },
     {
       "epoch": 0.29,
-      "grad_norm": 4.820233345031738,
-      "learning_rate": 8.50657683528373e-05,
-      "loss": 2.1814524841308596,
+      "grad_norm": 0.8435996174812317,
+      "learning_rate": 0.0003402630734113492,
+      "loss": 0.6225329971313477,
       "step": 8700
     },
     {
       "epoch": 0.29333333333333333,
-      "grad_norm": 4.834990501403809,
-      "learning_rate": 8.467075330211474e-05,
-      "loss": 2.14441162109375,
+      "grad_norm": 0.9413688778877258,
+      "learning_rate": 0.00033868301320845895,
+      "loss": 0.6248882293701172,
       "step": 8800
     },
     {
       "epoch": 0.2966666666666667,
-      "grad_norm": 4.696875095367432,
-      "learning_rate": 8.42715254701735e-05,
-      "loss": 2.1534934997558595,
+      "grad_norm": 0.7998372316360474,
+      "learning_rate": 0.000337086101880694,
+      "loss": 0.6171749496459961,
       "step": 8900
     },
     {
       "epoch": 0.3,
-      "grad_norm": 4.435513973236084,
-      "learning_rate": 8.386813336647339e-05,
-      "loss": 2.1541061401367188,
+      "grad_norm": 0.8585835099220276,
+      "learning_rate": 0.00033547253346589355,
+      "loss": 0.6206340789794922,
       "step": 9000
     },
     {
       "epoch": 0.30333333333333334,
-      "grad_norm": 4.625392436981201,
-      "learning_rate": 8.346062600646738e-05,
-      "loss": 2.167764892578125,
+      "grad_norm": 0.8483954071998596,
+      "learning_rate": 0.00033384250402586954,
+      "loss": 0.6195469665527343,
       "step": 9100
     },
     {
       "epoch": 0.30666666666666664,
-      "grad_norm": 4.518789291381836,
-      "learning_rate": 8.304905290564586e-05,
-      "loss": 2.125498504638672,
+      "grad_norm": 0.9016720056533813,
+      "learning_rate": 0.00033219621162258343,
+      "loss": 0.6200881958007812,
       "step": 9200
     },
     {
       "epoch": 0.31,
-      "grad_norm": 4.772843360900879,
-      "learning_rate": 8.263346407352018e-05,
-      "loss": 2.104402770996094,
+      "grad_norm": 0.829871416091919,
+      "learning_rate": 0.0003305338562940807,
+      "loss": 0.6127318954467773,
       "step": 9300
     },
     {
       "epoch": 1.002,
-      "grad_norm": 5.314713478088379,
-      "learning_rate": 8.22139100075459e-05,
-      "loss": 2.1437747192382814,
+      "grad_norm": 0.9081104397773743,
+      "learning_rate": 0.0003288556400301836,
+      "loss": 0.603317756652832,
       "step": 9400
     },
     {
       "epoch": 1.0053333333333334,
-      "grad_norm": 5.156741619110107,
-      "learning_rate": 8.179044168698721e-05,
-      "loss": 2.150110626220703,
+      "grad_norm": 0.8421012759208679,
+      "learning_rate": 0.00032716176674794884,
+      "loss": 0.5797752380371094,
       "step": 9500
     },
     {
       "epoch": 1.0086666666666666,
-      "grad_norm": 4.473363876342773,
-      "learning_rate": 8.136311056672223e-05,
-      "loss": 2.1256982421875,
+      "grad_norm": 0.8119559288024902,
+      "learning_rate": 0.0003254524422668889,
+      "loss": 0.5843183898925781,
       "step": 9600
     },
     {
       "epoch": 1.012,
-      "grad_norm": 4.629226207733154,
-      "learning_rate": 8.093196857099106e-05,
-      "loss": 2.0655870056152343,
+      "grad_norm": 0.8360543847084045,
+      "learning_rate": 0.00032372787428396423,
+      "loss": 0.5793385696411133,
       "step": 9700
     },
     {
       "epoch": 1.0153333333333334,
-      "grad_norm": 4.722533226013184,
-      "learning_rate": 8.04970680870864e-05,
-      "loss": 2.046549987792969,
+      "grad_norm": 0.780649721622467,
+      "learning_rate": 0.0003219882723483456,
+      "loss": 0.5682408142089844,
       "step": 9800
     },
     {
       "epoch": 1.0186666666666666,
-      "grad_norm": 4.950850486755371,
-      "learning_rate": 8.005846195898814e-05,
-      "loss": 2.065498809814453,
+      "grad_norm": 0.7468474507331848,
+      "learning_rate": 0.00032023384783595255,
+      "loss": 0.5659066772460938,
       "step": 9900
     },
     {
       "epoch": 1.022,
-      "grad_norm": 4.620322227478027,
-      "learning_rate": 7.961620348094241e-05,
-      "loss": 2.052088165283203,
+      "grad_norm": 0.8460882902145386,
+      "learning_rate": 0.00031846481392376966,
+      "loss": 0.5553397369384766,
       "step": 10000
     },
     {
       "epoch": 1.022,
-      "eval_loss": 2.195895195007324,
-      "eval_masked_accuracy": 0.48914785142024764,
-      "eval_runtime": 5.9829,
-      "eval_samples_per_second": 684.616,
-      "eval_steps_per_second": 2.674,
+      "eval_loss": 0.5976771116256714,
+      "eval_masked_accuracy": 0.8167562724014337,
+      "eval_runtime": 2.8203,
+      "eval_samples_per_second": 1452.351,
+      "eval_steps_per_second": 5.673,
       "step": 10000
     },
     {
       "epoch": 1.0253333333333334,
-      "grad_norm": 5.006026744842529,
-      "learning_rate": 7.917034639098579e-05,
-      "loss": 2.0433953857421874,
+      "grad_norm": 0.7821327447891235,
+      "learning_rate": 0.00031668138556394316,
+      "loss": 0.5550868606567383,
       "step": 10100
     },
     {
       "epoch": 1.0286666666666666,
-      "grad_norm": 4.821838855743408,
-      "learning_rate": 7.872094486441579e-05,
-      "loss": 2.0248480224609375,
+      "grad_norm": 1.0433557033538818,
+      "learning_rate": 0.00031488377945766314,
+      "loss": 0.5548404312133789,
       "step": 10200
     },
     {
       "epoch": 1.032,
-      "grad_norm": 5.129427433013916,
-      "learning_rate": 7.826805350720807e-05,
-      "loss": 2.034512481689453,
+      "grad_norm": 0.8580315113067627,
+      "learning_rate": 0.00031307221402883227,
+      "loss": 0.5406364440917969,
       "step": 10300
     },
     {
       "epoch": 1.0353333333333334,
-      "grad_norm": 4.920040130615234,
-      "learning_rate": 7.781172734938136e-05,
-      "loss": 2.0023085021972657,
+      "grad_norm": 0.8950467705726624,
+      "learning_rate": 0.00031124690939752543,
+      "loss": 0.533514404296875,
       "step": 10400
     },
     {
       "epoch": 1.0386666666666666,
-      "grad_norm": 4.989480495452881,
-      "learning_rate": 7.735202183831085e-05,
-      "loss": 2.0228404235839843,
+      "grad_norm": 0.7847508192062378,
+      "learning_rate": 0.0003094080873532434,
+      "loss": 0.538754539489746,
       "step": 10500
     },
     {
       "epoch": 1.042,
-      "grad_norm": 4.8022308349609375,
-      "learning_rate": 7.688899283199096e-05,
-      "loss": 2.004698944091797,
+      "grad_norm": 0.868969738483429,
+      "learning_rate": 0.00030755597132796385,
+      "loss": 0.5428607177734375,
       "step": 10600
     },
     {
       "epoch": 1.0453333333333332,
-      "grad_norm": 4.959015846252441,
-      "learning_rate": 7.642269659224803e-05,
-      "loss": 1.984919891357422,
+      "grad_norm": 0.7972626090049744,
+      "learning_rate": 0.0003056907863689921,
+      "loss": 0.5347858810424805,
       "step": 10700
     },
     {
       "epoch": 1.0486666666666666,
-      "grad_norm": 5.1211018562316895,
-      "learning_rate": 7.595318977790408e-05,
-      "loss": 1.9975146484375,
+      "grad_norm": 0.763969361782074,
+      "learning_rate": 0.0003038127591116163,
+      "loss": 0.532847785949707,
       "step": 10800
     },
     {
       "epoch": 1.052,
-      "grad_norm": 5.355852127075195,
-      "learning_rate": 7.54805294378923e-05,
-      "loss": 2.001761627197266,
+      "grad_norm": 0.7701809406280518,
+      "learning_rate": 0.0003019221177515692,
+      "loss": 0.5308395767211914,
       "step": 10900
     },
     {
       "epoch": 1.0553333333333332,
-      "grad_norm": 5.076831340789795,
-      "learning_rate": 7.50047730043252e-05,
-      "loss": 1.9820744323730468,
+      "grad_norm": 0.8959071636199951,
+      "learning_rate": 0.0003000190920173008,
+      "loss": 0.5308527374267578,
       "step": 11000
     },
     {
       "epoch": 1.0586666666666666,
-      "grad_norm": 5.53018045425415,
-      "learning_rate": 7.452597828551604e-05,
-      "loss": 1.9947706604003905,
+      "grad_norm": 1.0272624492645264,
+      "learning_rate": 0.00029810391314206416,
+      "loss": 0.5359263992309571,
       "step": 11100
     },
     {
       "epoch": 1.062,
-      "grad_norm": 5.11814546585083,
-      "learning_rate": 7.404420345895467e-05,
-      "loss": 1.990592498779297,
+      "grad_norm": 0.8793138861656189,
+      "learning_rate": 0.00029617681383581866,
+      "loss": 0.5368635559082031,
       "step": 11200
     },
     {
       "epoch": 1.0653333333333332,
-      "grad_norm": 5.014266490936279,
-      "learning_rate": 7.355950706423853e-05,
-      "loss": 1.9744894409179687,
+      "grad_norm": 0.9605848789215088,
+      "learning_rate": 0.0002942380282569541,
+      "loss": 0.5313293838500976,
       "step": 11300
     },
     {
       "epoch": 1.0686666666666667,
-      "grad_norm": 5.087450981140137,
-      "learning_rate": 7.307194799595958e-05,
-      "loss": 1.9660249328613282,
+      "grad_norm": 0.8334460258483887,
+      "learning_rate": 0.00029228779198383834,
+      "loss": 0.5360410690307618,
       "step": 11400
     },
     {
       "epoch": 1.072,
-      "grad_norm": 4.689304828643799,
-      "learning_rate": 7.25815854965481e-05,
-      "loss": 1.9629814147949218,
+      "grad_norm": 0.8171955347061157,
+      "learning_rate": 0.0002903263419861924,
+      "loss": 0.5353861999511719,
       "step": 11500
     },
     {
       "epoch": 1.0753333333333333,
-      "grad_norm": 4.9676008224487305,
-      "learning_rate": 7.20884791490743e-05,
-      "loss": 1.9364776611328125,
+      "grad_norm": 0.8680943250656128,
+      "learning_rate": 0.0002883539165962972,
+      "loss": 0.5278213119506836,
       "step": 11600
     },
     {
       "epoch": 1.0786666666666667,
-      "grad_norm": 5.643824100494385,
-      "learning_rate": 7.159268887000846e-05,
-      "loss": 1.93596435546875,
+      "grad_norm": 0.9106718897819519,
+      "learning_rate": 0.0002863707554800338,
+      "loss": 0.5310506439208984,
       "step": 11700
     },
     {
       "epoch": 1.082,
-      "grad_norm": 5.153912544250488,
-      "learning_rate": 7.109427490194056e-05,
-      "loss": 1.9413229370117187,
+      "grad_norm": 0.9140894412994385,
+      "learning_rate": 0.00028437709960776224,
+      "loss": 0.5164321899414063,
       "step": 11800
     },
     {
       "epoch": 1.0853333333333333,
-      "grad_norm": 4.910811424255371,
-      "learning_rate": 7.059329780626034e-05,
-      "loss": 1.9431932067871094,
+      "grad_norm": 0.790691077709198,
+      "learning_rate": 0.0002823731912250414,
+      "loss": 0.5228651428222656,
       "step": 11900
     },
     {
       "epoch": 1.0886666666666667,
-      "grad_norm": 5.167984962463379,
-      "learning_rate": 7.008981845579864e-05,
-      "loss": 1.9339044189453125,
+      "grad_norm": 0.757777988910675,
+      "learning_rate": 0.0002803592738231946,
+      "loss": 0.5215324401855469,
       "step": 12000
     },
     {
       "epoch": 1.092,
-      "grad_norm": 5.159730911254883,
-      "learning_rate": 6.958389802743077e-05,
-      "loss": 1.928137969970703,
+      "grad_norm": 0.8369808793067932,
+      "learning_rate": 0.00027833559210972307,
+      "loss": 0.5239771270751953,
       "step": 12100
     },
     {
       "epoch": 1.0953333333333333,
-      "grad_norm": 5.2495341300964355,
-      "learning_rate": 6.9075597994643e-05,
-      "loss": 1.9172648620605468,
+      "grad_norm": 0.8947746157646179,
+      "learning_rate": 0.000276302391978572,
+      "loss": 0.5179909896850586,
       "step": 12200
     },
     {
       "epoch": 1.0986666666666667,
-      "grad_norm": 5.271884918212891,
-      "learning_rate": 6.856498012006318e-05,
-      "loss": 1.9498109436035156,
+      "grad_norm": 0.9140170812606812,
+      "learning_rate": 0.00027425992048025273,
+      "loss": 0.5237076568603516,
       "step": 12300
     },
     {
       "epoch": 1.102,
-      "grad_norm": 4.869717121124268,
-      "learning_rate": 6.805210644795588e-05,
-      "loss": 1.9088368225097656,
+      "grad_norm": 0.8575257062911987,
+      "learning_rate": 0.00027220842579182353,
+      "loss": 0.5152669525146485,
       "step": 12400
     },
     {
       "epoch": 1.1053333333333333,
-      "grad_norm": 4.915546894073486,
-      "learning_rate": 6.753703929668362e-05,
-      "loss": 1.9135501098632812,
+      "grad_norm": 0.7598596215248108,
+      "learning_rate": 0.0002701481571867345,
+      "loss": 0.5134631729125977,
       "step": 12500
     },
     {
       "epoch": 1.1086666666666667,
-      "grad_norm": 4.548986434936523,
-      "learning_rate": 6.70198412511347e-05,
-      "loss": 1.9082426452636718,
+      "grad_norm": 0.8437865972518921,
+      "learning_rate": 0.0002680793650045388,
+      "loss": 0.5225432586669921,
       "step": 12600
     },
     {
       "epoch": 1.112,
-      "grad_norm": 5.1761579513549805,
-      "learning_rate": 6.650057515511848e-05,
-      "loss": 1.9233914184570313,
+      "grad_norm": 0.7969825863838196,
+      "learning_rate": 0.00026600230062047393,
+      "loss": 0.5257926177978516,
       "step": 12700
     },
     {
       "epoch": 1.1153333333333333,
-      "grad_norm": 4.711123943328857,
-      "learning_rate": 6.597930410372954e-05,
-      "loss": 1.9090065002441405,
+      "grad_norm": 0.8563022613525391,
+      "learning_rate": 0.00026391721641491817,
+      "loss": 0.5251951217651367,
       "step": 12800
     },
     {
       "epoch": 1.1186666666666667,
-      "grad_norm": 4.753604888916016,
-      "learning_rate": 6.54560914356809e-05,
-      "loss": 1.913031768798828,
+      "grad_norm": 0.708271861076355,
+      "learning_rate": 0.0002618243657427236,
+      "loss": 0.5297073745727539,
       "step": 12900
     },
     {
       "epoch": 1.1219999999999999,
-      "grad_norm": 4.507719993591309,
-      "learning_rate": 6.493100072560799e-05,
-      "loss": 1.9241600036621094,
+      "grad_norm": 0.7699816823005676,
+      "learning_rate": 0.00025972400290243195,
+      "loss": 0.5333083724975586,
       "step": 13000
     },
     {
       "epoch": 1.1253333333333333,
-      "grad_norm": 4.860040187835693,
-      "learning_rate": 6.44040957763438e-05,
-      "loss": 1.921043701171875,
+      "grad_norm": 0.8275418877601624,
+      "learning_rate": 0.0002576163831053752,
+      "loss": 0.5334249114990235,
       "step": 13100
     },
     {
       "epoch": 1.1286666666666667,
-      "grad_norm": 4.9244465827941895,
-      "learning_rate": 6.387544061116621e-05,
-      "loss": 1.9413529968261718,
+      "grad_norm": 0.83842533826828,
+      "learning_rate": 0.00025550176244466485,
+      "loss": 0.5315378189086915,
       "step": 13200
     },
     {
       "epoch": 1.1320000000000001,
-      "grad_norm": 4.874339580535889,
-      "learning_rate": 6.334509946601879e-05,
-      "loss": 1.9148121643066407,
+      "grad_norm": 0.8859015703201294,
+      "learning_rate": 0.00025338039786407517,
+      "loss": 0.5462711715698242,
       "step": 13300
     },
     {
       "epoch": 1.1353333333333333,
-      "grad_norm": 4.540445804595947,
-      "learning_rate": 6.281313678170542e-05,
-      "loss": 1.923395538330078,
+      "grad_norm": 0.7318875193595886,
+      "learning_rate": 0.0002512525471268217,
+      "loss": 0.5278379440307617,
       "step": 13400
     },
     {
       "epoch": 1.1386666666666667,
-      "grad_norm": 5.011052131652832,
-      "learning_rate": 6.227961719606039e-05,
-      "loss": 1.9019638061523438,
+      "grad_norm": 0.7980539798736572,
+      "learning_rate": 0.00024911846878424156,
+      "loss": 0.5409846115112305,
       "step": 13500
     },
     {
       "epoch": 1.142,
-      "grad_norm": 5.090750694274902,
-      "learning_rate": 6.174460553609426e-05,
-      "loss": 1.9019752502441407,
+      "grad_norm": 0.8805834650993347,
+      "learning_rate": 0.00024697842214437704,
+      "loss": 0.538138427734375,
       "step": 13600
     },
     {
       "epoch": 1.1453333333333333,
-      "grad_norm": 5.096031665802002,
-      "learning_rate": 6.120816681011684e-05,
-      "loss": 1.9170556640625,
+      "grad_norm": 0.8037357330322266,
+      "learning_rate": 0.00024483266724046736,
+      "loss": 0.528113784790039,
       "step": 13700
     },
     {
       "epoch": 1.1486666666666667,
-      "grad_norm": 5.055447101593018,
-      "learning_rate": 6.0670366199838215e-05,
-      "loss": 1.8996456909179686,
+      "grad_norm": 0.8429805636405945,
+      "learning_rate": 0.00024268146479935286,
+      "loss": 0.5357782363891601,
       "step": 13800
     },
     {
       "epoch": 1.152,
-      "grad_norm": 4.6030097007751465,
-      "learning_rate": 6.013126905244858e-05,
-      "loss": 1.9002468872070313,
+      "grad_norm": 0.876872181892395,
+      "learning_rate": 0.0002405250762097943,
+      "loss": 0.5338603591918946,
       "step": 13900
     },
     {
       "epoch": 1.1553333333333333,
-      "grad_norm": 4.697137355804443,
-      "learning_rate": 5.959094087267804e-05,
-      "loss": 1.886578826904297,
+      "grad_norm": 0.798278272151947,
+      "learning_rate": 0.00023836376349071215,
+      "loss": 0.5351076889038086,
       "step": 14000
     },
     {
       "epoch": 1.1586666666666667,
-      "grad_norm": 4.779449462890625,
-      "learning_rate": 5.904944731483723e-05,
-      "loss": 1.8922029113769532,
+      "grad_norm": 0.8131340742111206,
+      "learning_rate": 0.00023619778925934893,
+      "loss": 0.5315785598754883,
       "step": 14100
     },
     {
       "epoch": 1.162,
-      "grad_norm": 4.6244893074035645,
-      "learning_rate": 5.850685417483983e-05,
-      "loss": 1.894721221923828,
+      "grad_norm": 0.8453519940376282,
+      "learning_rate": 0.00023402741669935933,
+      "loss": 0.5312258911132812,
       "step": 14200
     },
     {
       "epoch": 1.1653333333333333,
-      "grad_norm": 4.580394268035889,
-      "learning_rate": 5.796322738220774e-05,
-      "loss": 1.8998484802246094,
+      "grad_norm": 0.7479920387268066,
+      "learning_rate": 0.00023185290952883097,
+      "loss": 0.5305691909790039,
       "step": 14300
     },
     {
       "epoch": 1.1686666666666667,
-      "grad_norm": 4.821572780609131,
-      "learning_rate": 5.741863299206014e-05,
-      "loss": 1.881027069091797,
+      "grad_norm": 0.7968600392341614,
+      "learning_rate": 0.00022967453196824056,
+      "loss": 0.5351016998291016,
       "step": 14400
     },
     {
       "epoch": 1.172,
-      "grad_norm": 4.807371139526367,
-      "learning_rate": 5.687313717708728e-05,
-      "loss": 1.8675735473632813,
+      "grad_norm": 0.7319141030311584,
+      "learning_rate": 0.0002274925487083491,
+      "loss": 0.523140754699707,
       "step": 14500
     },
     {
       "epoch": 1.1753333333333333,
-      "grad_norm": 4.700822830200195,
-      "learning_rate": 5.632680621950992e-05,
-      "loss": 1.8618142700195313,
+      "grad_norm": 0.7255148887634277,
+      "learning_rate": 0.00022530722487803967,
+      "loss": 0.5274886322021485,
       "step": 14600
     },
     {
       "epoch": 1.1786666666666668,
-      "grad_norm": 4.157495021820068,
-      "learning_rate": 5.57797065030255e-05,
-      "loss": 1.8685194396972655,
+      "grad_norm": 0.8094624876976013,
+      "learning_rate": 0.000223118826012102,
+      "loss": 0.5305857467651367,
       "step": 14700
     },
     {
       "epoch": 1.182,
-      "grad_norm": 4.889394283294678,
-      "learning_rate": 5.523190450474197e-05,
-      "loss": 1.886773681640625,
+      "grad_norm": 0.7611061930656433,
+      "learning_rate": 0.0002209276180189679,
+      "loss": 0.5199204635620117,
       "step": 14800
     },
     {
       "epoch": 1.1853333333333333,
-      "grad_norm": 4.845384120941162,
-      "learning_rate": 5.4683466787100304e-05,
-      "loss": 1.8534288024902343,
+      "grad_norm": 0.7731903195381165,
+      "learning_rate": 0.00021873386714840122,
+      "loss": 0.5233501434326172,
       "step": 14900
     },
     {
       "epoch": 1.1886666666666668,
-      "grad_norm": 4.679961204528809,
-      "learning_rate": 5.413445998978658e-05,
-      "loss": 1.8429244995117187,
+      "grad_norm": 0.802269697189331,
+      "learning_rate": 0.0002165378399591463,
+      "loss": 0.5173654556274414,
       "step": 15000
     },
     {
       "epoch": 1.1886666666666668,
-      "eval_loss": 1.8467531204223633,
-      "eval_masked_accuracy": 0.5551868417255807,
-      "eval_runtime": 6.0839,
-      "eval_samples_per_second": 673.248,
-      "eval_steps_per_second": 2.63,
+      "eval_loss": 0.4867754876613617,
+      "eval_masked_accuracy": 0.8465927303032246,
+      "eval_runtime": 2.9083,
+      "eval_samples_per_second": 1408.378,
+      "eval_steps_per_second": 5.501,
       "step": 15000
     },
     {
       "epoch": 1.192,
-      "grad_norm": 4.631464004516602,
-      "learning_rate": 5.358495082163476e-05,
-      "loss": 1.8469007873535157,
+      "grad_norm": 0.8154687285423279,
+      "learning_rate": 0.00021433980328653903,
+      "loss": 0.5211496353149414,
       "step": 15100
     },
     {
       "epoch": 1.1953333333333334,
-      "grad_norm": 4.638521194458008,
-      "learning_rate": 5.3035006052520955e-05,
-      "loss": 1.8356764221191406,
+      "grad_norm": 0.870178759098053,
+      "learning_rate": 0.00021214002421008382,
+      "loss": 0.5156788253784179,
       "step": 15200
     },
     {
       "epoch": 1.1986666666666665,
-      "grad_norm": 4.570642471313477,
-      "learning_rate": 5.2484692505250376e-05,
-      "loss": 1.8426182556152344,
+      "grad_norm": 0.7339473962783813,
+      "learning_rate": 0.0002099387700210015,
+      "loss": 0.5113505935668945,
       "step": 15300
     },
     {
       "epoch": 1.202,
-      "grad_norm": 4.465674877166748,
-      "learning_rate": 5.193407704743782e-05,
-      "loss": 1.8424899291992187,
+      "grad_norm": 0.8333882689476013,
+      "learning_rate": 0.00020773630818975128,
+      "loss": 0.5066380310058594,
       "step": 15400
     },
     {
       "epoch": 1.2053333333333334,
-      "grad_norm": 4.427441596984863,
-      "learning_rate": 5.1383226583382684e-05,
-      "loss": 1.8333729553222655,
+      "grad_norm": 0.7248000502586365,
+      "learning_rate": 0.00020553290633353074,
+      "loss": 0.5148737716674805,
       "step": 15500
     },
     {
       "epoch": 1.2086666666666668,
-      "grad_norm": 4.944128036499023,
-      "learning_rate": 5.0832208045939556e-05,
-      "loss": 1.8115110778808594,
+      "grad_norm": 0.8077163696289062,
+      "learning_rate": 0.00020332883218375823,
+      "loss": 0.5026751327514648,
       "step": 15600
     },
     {
       "epoch": 1.212,
-      "grad_norm": 4.622330188751221,
-      "learning_rate": 5.028108838838533e-05,
-      "loss": 1.8149687194824218,
+      "grad_norm": 0.7529610991477966,
+      "learning_rate": 0.00020112435355354132,
+      "loss": 0.5074459457397461,
       "step": 15700
     },
     {
       "epoch": 1.2153333333333334,
-      "grad_norm": 4.573352336883545,
-      "learning_rate": 4.972993457628381e-05,
-      "loss": 1.8161116027832032,
+      "grad_norm": 0.8418951630592346,
+      "learning_rate": 0.00019891973830513525,
+      "loss": 0.4947340393066406,
       "step": 15800
     },
     {
       "epoch": 1.2186666666666666,
-      "grad_norm": 4.753305435180664,
-      "learning_rate": 4.9178813579348865e-05,
-      "loss": 1.8010633850097657,
+      "grad_norm": 0.7463741302490234,
+      "learning_rate": 0.00019671525431739546,
+      "loss": 0.49930683135986326,
       "step": 15900
     },
     {
       "epoch": 1.222,
-      "grad_norm": 5.017966270446777,
-      "learning_rate": 4.862779236330705e-05,
-      "loss": 1.8063246154785155,
+      "grad_norm": 0.7505474090576172,
+      "learning_rate": 0.0001945111694532282,
+      "loss": 0.49693927764892576,
       "step": 16000
     },
     {
       "epoch": 1.2253333333333334,
-      "grad_norm": 4.491023063659668,
-      "learning_rate": 4.807693788176071e-05,
-      "loss": 1.8086869812011719,
+      "grad_norm": 0.7937129735946655,
+      "learning_rate": 0.00019230775152704284,
+      "loss": 0.5052358627319335,
       "step": 16100
     },
     {
       "epoch": 1.2286666666666666,
-      "grad_norm": 4.622035503387451,
-      "learning_rate": 4.752631706805261e-05,
-      "loss": 1.7760255432128906,
+      "grad_norm": 0.6797370910644531,
+      "learning_rate": 0.00019010526827221045,
+      "loss": 0.4982032775878906,
       "step": 16200
     },
     {
       "epoch": 1.232,
-      "grad_norm": 4.742088317871094,
-      "learning_rate": 4.6975996827132926e-05,
-      "loss": 1.7725007629394531,
+      "grad_norm": 0.755185604095459,
+      "learning_rate": 0.0001879039873085317,
+      "loss": 0.4924372482299805,
       "step": 16300
     },
     {
       "epoch": 1.2353333333333334,
-      "grad_norm": 4.786332130432129,
-      "learning_rate": 4.6426044027429786e-05,
-      "loss": 1.7803855895996095,
+      "grad_norm": 0.6874153017997742,
+      "learning_rate": 0.00018570417610971914,
+      "loss": 0.49145896911621095,
       "step": 16400
     },
     {
       "epoch": 1.2386666666666666,
-      "grad_norm": 4.642701625823975,
-      "learning_rate": 4.58765254927242e-05,
-      "loss": 1.7837786865234375,
+      "grad_norm": 0.785365879535675,
+      "learning_rate": 0.0001835061019708968,
+      "loss": 0.4896976470947266,
       "step": 16500
     },
     {
       "epoch": 1.242,
-      "grad_norm": 4.578219890594482,
-      "learning_rate": 4.5327507994030394e-05,
-      "loss": 1.7726612854003907,
+      "grad_norm": 0.7931706309318542,
+      "learning_rate": 0.00018131003197612158,
+      "loss": 0.48486351013183593,
       "step": 16600
     },
     {
       "epoch": 1.2453333333333334,
-      "grad_norm": 4.572558403015137,
-      "learning_rate": 4.47790582414826e-05,
-      "loss": 1.7728559875488281,
+      "grad_norm": 0.711540162563324,
+      "learning_rate": 0.0001791162329659304,
+      "loss": 0.4865373992919922,
       "step": 16700
     },
     {
       "epoch": 1.2486666666666666,
-      "grad_norm": 4.7602434158325195,
-      "learning_rate": 4.423124287622926e-05,
-      "loss": 1.7534779357910155,
+      "grad_norm": 0.8093428611755371,
+      "learning_rate": 0.00017692497150491703,
+      "loss": 0.4822272491455078,
       "step": 16800
     },
     {
       "epoch": 1.252,
-      "grad_norm": 4.802716255187988,
-      "learning_rate": 4.3684128462335533e-05,
-      "loss": 1.7711273193359376,
+      "grad_norm": 0.8606541156768799,
+      "learning_rate": 0.00017473651384934213,
+      "loss": 0.47816139221191406,
       "step": 16900
     },
     {
       "epoch": 1.2553333333333334,
-      "grad_norm": 4.785411357879639,
-      "learning_rate": 4.3137781478695236e-05,
-      "loss": 1.752899169921875,
+      "grad_norm": 0.6578624248504639,
+      "learning_rate": 0.00017255112591478094,
+      "loss": 0.4819015121459961,
       "step": 17000
     },
     {
       "epoch": 1.2586666666666666,
-      "grad_norm": 4.973933219909668,
-      "learning_rate": 4.25922683109531e-05,
-      "loss": 1.7573219299316407,
+      "grad_norm": 0.8381503224372864,
+      "learning_rate": 0.0001703690732438124,
+      "loss": 0.4800693893432617,
       "step": 17100
     },
     {
       "epoch": 1.262,
-      "grad_norm": 4.814939022064209,
-      "learning_rate": 4.204765524343841e-05,
-      "loss": 1.7253883361816407,
+      "grad_norm": 0.9081389904022217,
+      "learning_rate": 0.00016819062097375365,
+      "loss": 0.4796139907836914,
       "step": 17200
     },
     {
       "epoch": 1.2653333333333334,
-      "grad_norm": 4.742300987243652,
-      "learning_rate": 4.1504008451110845e-05,
-      "loss": 1.7295744323730469,
+      "grad_norm": 0.8076674938201904,
+      "learning_rate": 0.00016601603380444338,
+      "loss": 0.47347488403320315,
       "step": 17300
     },
     {
       "epoch": 1.2686666666666666,
-      "grad_norm": 4.9945969581604,
-      "learning_rate": 4.096139399151971e-05,
-      "loss": 1.7228933715820312,
+      "grad_norm": 0.7531787157058716,
+      "learning_rate": 0.00016384557596607884,
+      "loss": 0.4700050354003906,
       "step": 17400
     },
     {
       "epoch": 1.272,
-      "grad_norm": 4.486416339874268,
-      "learning_rate": 4.041987779677745e-05,
-      "loss": 1.7148330688476563,
+      "grad_norm": 0.843329131603241,
+      "learning_rate": 0.0001616795111871098,
+      "loss": 0.47036163330078123,
       "step": 17500
     },
     {
       "epoch": 1.2753333333333332,
-      "grad_norm": 4.558888912200928,
-      "learning_rate": 3.9879525665548276e-05,
-      "loss": 1.7216806030273437,
+      "grad_norm": 0.7025447487831116,
+      "learning_rate": 0.0001595181026621931,
+      "loss": 0.46735321044921874,
       "step": 17600
     },
     {
       "epoch": 1.2786666666666666,
-      "grad_norm": 4.67527437210083,
-      "learning_rate": 3.9340403255053135e-05,
-      "loss": 1.700666961669922,
+      "grad_norm": 0.7623017430305481,
+      "learning_rate": 0.00015736161302021254,
+      "loss": 0.46535400390625,
       "step": 17700
     },
     {
       "epoch": 1.282,
-      "grad_norm": 4.606650352478027,
-      "learning_rate": 3.880257607309178e-05,
-      "loss": 1.7064649963378906,
+      "grad_norm": 0.7620524764060974,
+      "learning_rate": 0.00015521030429236713,
+      "loss": 0.4652018356323242,
       "step": 17800
     },
     {
       "epoch": 1.2853333333333334,
-      "grad_norm": 4.782898902893066,
-      "learning_rate": 3.826610947008312e-05,
-      "loss": 1.6886534118652343,
+      "grad_norm": 0.7925307154655457,
+      "learning_rate": 0.00015306443788033248,
+      "loss": 0.46518348693847655,
       "step": 17900
     },
     {
       "epoch": 1.2886666666666666,
-      "grad_norm": 4.922590732574463,
-      "learning_rate": 3.773106863112451e-05,
-      "loss": 1.703627471923828,
+      "grad_norm": 0.7878969311714172,
+      "learning_rate": 0.00015092427452449803,
+      "loss": 0.4622929382324219,
       "step": 18000
     },
     {
       "epoch": 1.292,
-      "grad_norm": 4.657540321350098,
-      "learning_rate": 3.719751856807125e-05,
-      "loss": 1.6833297729492187,
+      "grad_norm": 0.7757784724235535,
+      "learning_rate": 0.000148790074272285,
+      "loss": 0.4650875091552734,
       "step": 18100
     },
     {
       "epoch": 1.2953333333333332,
-      "grad_norm": 4.772761344909668,
-      "learning_rate": 3.666552411163718e-05,
-      "loss": 1.688351593017578,
+      "grad_norm": 0.735175371170044,
+      "learning_rate": 0.00014666209644654873,
+      "loss": 0.4549406433105469,
       "step": 18200
     },
     {
       "epoch": 1.2986666666666666,
-      "grad_norm": 4.486453056335449,
-      "learning_rate": 3.6135149903517115e-05,
-      "loss": 1.6901985168457032,
+      "grad_norm": 0.6873441934585571,
+      "learning_rate": 0.00014454059961406846,
+      "loss": 0.4520915985107422,
       "step": 18300
     },
     {
       "epoch": 1.302,
-      "grad_norm": 4.456425189971924,
-      "learning_rate": 3.56064603885324e-05,
-      "loss": 1.6628305053710937,
+      "grad_norm": 0.7766420841217041,
+      "learning_rate": 0.0001424258415541296,
+      "loss": 0.447607536315918,
       "step": 18400
     },
     {
       "epoch": 1.3053333333333335,
-      "grad_norm": 4.608875751495361,
-      "learning_rate": 3.5079519806800374e-05,
-      "loss": 1.6773171997070313,
+      "grad_norm": 0.6834190487861633,
+      "learning_rate": 0.0001403180792272015,
+      "loss": 0.4524835205078125,
       "step": 18500
     },
     {
       "epoch": 1.3086666666666666,
-      "grad_norm": 4.784965515136719,
-      "learning_rate": 3.4554392185928564e-05,
-      "loss": 1.6454412841796875,
+      "grad_norm": 0.7714037299156189,
+      "learning_rate": 0.00013821756874371426,
+      "loss": 0.43865036010742187,
       "step": 18600
     },
     {
       "epoch": 2.0006666666666666,
-      "grad_norm": 4.08699893951416,
-      "learning_rate": 3.40311413332349e-05,
-      "loss": 1.6564366149902343,
+      "grad_norm": 0.6421045064926147,
+      "learning_rate": 0.0001361245653329396,
+      "loss": 0.45037097930908204,
       "step": 18700
     },
     {
       "epoch": 2.004,
-      "grad_norm": 4.481525897979736,
-      "learning_rate": 3.350983082799456e-05,
-      "loss": 1.62840576171875,
+      "grad_norm": 0.8619283437728882,
+      "learning_rate": 0.00013403932331197823,
+      "loss": 0.4485680389404297,
       "step": 18800
     },
     {
       "epoch": 2.0073333333333334,
-      "grad_norm": 3.9877312183380127,
-      "learning_rate": 3.299052401371456e-05,
-      "loss": 1.5869935607910157,
+      "grad_norm": 0.7336652278900146,
+      "learning_rate": 0.00013196209605485825,
+      "loss": 0.4345798110961914,
       "step": 18900
     },
     {
       "epoch": 2.010666666666667,
-      "grad_norm": 4.130960464477539,
-      "learning_rate": 3.247328399043706e-05,
-      "loss": 1.570656280517578,
+      "grad_norm": 0.7360154986381531,
+      "learning_rate": 0.00012989313596174825,
+      "loss": 0.4353963470458984,
       "step": 19000
     },
     {
       "epoch": 2.014,
-      "grad_norm": 4.480629920959473,
-      "learning_rate": 3.195817360707207e-05,
-      "loss": 1.5836396789550782,
+      "grad_norm": 0.7875688076019287,
+      "learning_rate": 0.0001278326944282883,
+      "loss": 0.4323548126220703,
       "step": 19100
     },
     {
       "epoch": 2.017333333333333,
-      "grad_norm": 4.205255031585693,
-      "learning_rate": 3.144525545376095e-05,
-      "loss": 1.5562339782714845,
+      "grad_norm": 0.8237797617912292,
+      "learning_rate": 0.0001257810218150438,
+      "loss": 0.43446548461914064,
       "step": 19200
     },
     {
       "epoch": 2.0206666666666666,
-      "grad_norm": 4.235847473144531,
-      "learning_rate": 3.093459185427102e-05,
-      "loss": 1.5763328552246094,
+      "grad_norm": 0.775242030620575,
+      "learning_rate": 0.0001237383674170841,
+      "loss": 0.4317903900146484,
       "step": 19300
     },
     {
       "epoch": 2.024,
-      "grad_norm": 4.205441951751709,
-      "learning_rate": 3.0426244858422847e-05,
-      "loss": 1.5626734924316406,
+      "grad_norm": 0.7316193580627441,
+      "learning_rate": 0.00012170497943369139,
+      "loss": 0.4266176223754883,
       "step": 19400
     },
     {
       "epoch": 2.0273333333333334,
-      "grad_norm": 4.186418533325195,
-      "learning_rate": 2.9920276234550636e-05,
-      "loss": 1.575146484375,
+      "grad_norm": 0.8828946352005005,
+      "learning_rate": 0.00011968110493820255,
+      "loss": 0.42110267639160154,
       "step": 19500
     },
     {
       "epoch": 2.030666666666667,
-      "grad_norm": 4.292003154754639,
-      "learning_rate": 2.9416747461996853e-05,
-      "loss": 1.54133544921875,
+      "grad_norm": 0.7441554665565491,
+      "learning_rate": 0.00011766698984798741,
+      "loss": 0.42277420043945313,
       "step": 19600
     },
     {
       "epoch": 2.034,
-      "grad_norm": 4.439959526062012,
-      "learning_rate": 2.8915719723641975e-05,
-      "loss": 1.5504551696777344,
+      "grad_norm": 0.7916893362998962,
+      "learning_rate": 0.0001156628788945679,
+      "loss": 0.42497276306152343,
       "step": 19700
     },
     {
       "epoch": 2.037333333333333,
-      "grad_norm": 4.374598503112793,
-      "learning_rate": 2.841725389847032e-05,
-      "loss": 1.5403448486328124,
+      "grad_norm": 0.7108286023139954,
+      "learning_rate": 0.00011366901559388127,
+      "loss": 0.42226448059082033,
       "step": 19800
     },
     {
       "epoch": 2.0406666666666666,
-      "grad_norm": 4.581845760345459,
-      "learning_rate": 2.7921410554172723e-05,
-      "loss": 1.550257568359375,
+      "grad_norm": 0.7501102089881897,
+      "learning_rate": 0.00011168564221669089,
+      "loss": 0.4148076629638672,
       "step": 19900
     },
     {
       "epoch": 2.044,
-      "grad_norm": 4.183570384979248,
-      "learning_rate": 2.7428249939787e-05,
-      "loss": 1.5298226928710938,
+      "grad_norm": 0.7276322841644287,
+      "learning_rate": 0.000109712999759148,
+      "loss": 0.41933116912841795,
       "step": 20000
     },
     {
       "epoch": 2.044,
-      "eval_loss": 1.6956673860549927,
-      "eval_masked_accuracy": 0.5865110984632897,
-      "eval_runtime": 6.0932,
-      "eval_samples_per_second": 672.222,
-      "eval_steps_per_second": 2.626,
+      "eval_loss": 0.4309006333351135,
+      "eval_masked_accuracy": 0.862885533617241,
+      "eval_runtime": 2.9684,
+      "eval_samples_per_second": 1379.874,
+      "eval_steps_per_second": 5.39,
       "step": 20000
     },
     {
       "epoch": 2.0473333333333334,
-      "grad_norm": 4.26902437210083,
-      "learning_rate": 2.6937831978377325e-05,
-      "loss": 1.5412590026855468,
+      "grad_norm": 0.8200253844261169,
+      "learning_rate": 0.0001077513279135093,
+      "loss": 0.4121805572509766,
       "step": 20100
     },
     {
       "epoch": 2.050666666666667,
-      "grad_norm": 4.665940761566162,
-      "learning_rate": 2.6450216259753e-05,
-      "loss": 1.5438227844238281,
+      "grad_norm": 0.7574154734611511,
+      "learning_rate": 0.000105800865039012,
+      "loss": 0.4140956115722656,
       "step": 20200
     },
     {
       "epoch": 2.054,
-      "grad_norm": 4.335201740264893,
-      "learning_rate": 2.596546203322777e-05,
-      "loss": 1.5362275695800782,
+      "grad_norm": 0.7535119652748108,
+      "learning_rate": 0.00010386184813291108,
+      "loss": 0.4128216552734375,
       "step": 20300
     },
     {
       "epoch": 2.0573333333333332,
-      "grad_norm": 4.486615180969238,
-      "learning_rate": 2.5483628200420644e-05,
-      "loss": 1.53197021484375,
+      "grad_norm": 0.7665644884109497,
+      "learning_rate": 0.00010193451280168258,
+      "loss": 0.4141941833496094,
       "step": 20400
     },
     {
       "epoch": 2.0606666666666666,
-      "grad_norm": 4.637569904327393,
-      "learning_rate": 2.500477330809881e-05,
-      "loss": 1.510602264404297,
+      "grad_norm": 0.6833004951477051,
+      "learning_rate": 0.00010001909323239524,
+      "loss": 0.41669425964355467,
       "step": 20500
     },
     {
       "epoch": 2.064,
-      "grad_norm": 4.1064558029174805,
-      "learning_rate": 2.452895554106368e-05,
-      "loss": 1.502353973388672,
+      "grad_norm": 0.6854680180549622,
+      "learning_rate": 9.811582216425472e-05,
+      "loss": 0.40588130950927737,
       "step": 20600
     },
     {
       "epoch": 2.0673333333333335,
-      "grad_norm": 4.285604953765869,
-      "learning_rate": 2.405623271508108e-05,
-      "loss": 1.528131866455078,
+      "grad_norm": 0.7551723122596741,
+      "learning_rate": 9.622493086032432e-05,
+      "loss": 0.40616172790527344,
       "step": 20700
     },
     {
       "epoch": 2.070666666666667,
-      "grad_norm": 4.3571367263793945,
-      "learning_rate": 2.358666226985599e-05,
-      "loss": 1.530682373046875,
+      "grad_norm": 0.7875136733055115,
+      "learning_rate": 9.434664907942397e-05,
+      "loss": 0.40618175506591797,
       "step": 20800
     },
     {
       "epoch": 2.074,
-      "grad_norm": 4.226117134094238,
-      "learning_rate": 2.312030126205335e-05,
-      "loss": 1.5152372741699218,
+      "grad_norm": 0.7156426906585693,
+      "learning_rate": 9.24812050482134e-05,
+      "loss": 0.40087387084960935,
       "step": 20900
     },
     {
       "epoch": 2.0773333333333333,
-      "grad_norm": 4.25309419631958,
-      "learning_rate": 2.2657206358364997e-05,
-      "loss": 1.5110414123535156,
+      "grad_norm": 0.7035794258117676,
+      "learning_rate": 9.062882543345999e-05,
+      "loss": 0.3979297637939453,
       "step": 21000
     },
     {
       "epoch": 2.0806666666666667,
-      "grad_norm": 4.612072467803955,
-      "learning_rate": 2.219743382862437e-05,
-      "loss": 1.4938458251953124,
+      "grad_norm": 0.7434916496276855,
+      "learning_rate": 8.878973531449748e-05,
+      "loss": 0.3994139099121094,
       "step": 21100
     },
     {
       "epoch": 2.084,
-      "grad_norm": 4.316677093505859,
-      "learning_rate": 2.1741039538969183e-05,
-      "loss": 1.4943765258789063,
+      "grad_norm": 0.6868703961372375,
+      "learning_rate": 8.696415815587673e-05,
+      "loss": 0.39528240203857423,
       "step": 21200
     },
     {
       "epoch": 2.0873333333333335,
-      "grad_norm": 4.053011417388916,
-      "learning_rate": 2.1288078945053193e-05,
-      "loss": 1.486304931640625,
+      "grad_norm": 0.7844119071960449,
+      "learning_rate": 8.515231578021277e-05,
+      "loss": 0.3911682891845703,
       "step": 21300
     },
     {
       "epoch": 2.0906666666666665,
-      "grad_norm": 4.295980453491211,
-      "learning_rate": 2.083860708530798e-05,
-      "loss": 1.4813737487792968,
+      "grad_norm": 0.7842739820480347,
+      "learning_rate": 8.335442834123192e-05,
+      "loss": 0.3987908935546875,
       "step": 21400
     },
     {
       "epoch": 2.094,
-      "grad_norm": 4.288027763366699,
-      "learning_rate": 2.0392678574255282e-05,
-      "loss": 1.490546875,
+      "grad_norm": 0.7348713874816895,
+      "learning_rate": 8.157071429702113e-05,
+      "loss": 0.38934185028076174,
       "step": 21500
     },
     {
       "epoch": 2.0973333333333333,
-      "grad_norm": 4.3048505783081055,
-      "learning_rate": 1.99503475958708e-05,
-      "loss": 1.4724571228027343,
+      "grad_norm": 0.7054168581962585,
+      "learning_rate": 7.98013903834832e-05,
+      "loss": 0.39156261444091794,
       "step": 21600
     },
     {
       "epoch": 2.1006666666666667,
-      "grad_norm": 4.423007488250732,
-      "learning_rate": 1.9511667897000574e-05,
-      "loss": 1.475355682373047,
+      "grad_norm": 0.7588643431663513,
+      "learning_rate": 7.80466715880023e-05,
+      "loss": 0.38873050689697264,
       "step": 21700
     },
     {
       "epoch": 2.104,
-      "grad_norm": 4.3376336097717285,
-      "learning_rate": 1.9076692780830114e-05,
-      "loss": 1.4873141479492187,
+      "grad_norm": 0.7288216948509216,
+      "learning_rate": 7.630677112332046e-05,
+      "loss": 0.3866202926635742,
       "step": 21800
     },
     {
       "epoch": 2.1073333333333335,
-      "grad_norm": 4.22784423828125,
-      "learning_rate": 1.8645475100407788e-05,
-      "loss": 1.5029916381835937,
+      "grad_norm": 0.7503123879432678,
+      "learning_rate": 7.458190040163115e-05,
+      "loss": 0.39182708740234373,
       "step": 21900
     },
     {
       "epoch": 2.1106666666666665,
-      "grad_norm": 4.327232360839844,
-      "learning_rate": 1.8218067252222597e-05,
-      "loss": 1.5122482299804687,
+      "grad_norm": 0.7055880427360535,
+      "learning_rate": 7.287226900889039e-05,
+      "loss": 0.401009635925293,
       "step": 22000
     },
     {
       "epoch": 2.114,
-      "grad_norm": 4.530300140380859,
-      "learning_rate": 1.779452116983769e-05,
-      "loss": 1.5061280822753906,
+      "grad_norm": 0.8966184258460999,
+      "learning_rate": 7.117808467935076e-05,
+      "loss": 0.39603382110595703,
       "step": 22100
     },
     {
       "epoch": 2.1173333333333333,
-      "grad_norm": 4.228301525115967,
-      "learning_rate": 1.737488831757997e-05,
-      "loss": 1.52053955078125,
+      "grad_norm": 0.6970842480659485,
+      "learning_rate": 6.949955327031988e-05,
+      "loss": 0.398880729675293,
       "step": 22200
     },
     {
       "epoch": 2.1206666666666667,
-      "grad_norm": 4.379787921905518,
-      "learning_rate": 1.695921968428668e-05,
-      "loss": 1.5291818237304688,
+      "grad_norm": 0.7583353519439697,
+      "learning_rate": 6.783687873714672e-05,
+      "loss": 0.3997147369384766,
       "step": 22300
     },
     {
       "epoch": 2.124,
-      "grad_norm": 4.750154972076416,
-      "learning_rate": 1.6547565777109976e-05,
-      "loss": 1.5258726501464843,
+      "grad_norm": 0.7861075401306152,
+      "learning_rate": 6.61902631084399e-05,
+      "loss": 0.4028291702270508,
       "step": 22400
     },
     {
       "epoch": 2.1273333333333335,
-      "grad_norm": 5.20265007019043,
-      "learning_rate": 1.613997661537981e-05,
-      "loss": 1.5335383605957031,
+      "grad_norm": 0.6678594946861267,
+      "learning_rate": 6.455990646151924e-05,
+      "loss": 0.3987395095825195,
       "step": 22500
     },
     {
       "epoch": 2.1306666666666665,
-      "grad_norm": 4.6225433349609375,
-      "learning_rate": 1.573650172452615e-05,
-      "loss": 1.5420014953613281,
+      "grad_norm": 0.7038342952728271,
+      "learning_rate": 6.29460068981046e-05,
+      "loss": 0.39326480865478514,
       "step": 22600
     },
     {
       "epoch": 2.134,
-      "grad_norm": 4.721563339233398,
-      "learning_rate": 1.5337190130061318e-05,
-      "loss": 1.5496218872070313,
+      "grad_norm": 0.7392048239707947,
+      "learning_rate": 6.134876052024527e-05,
+      "loss": 0.40042667388916015,
       "step": 22700
     },
     {
       "epoch": 2.1373333333333333,
-      "grad_norm": 4.762903690338135,
-      "learning_rate": 1.4942090351622883e-05,
-      "loss": 1.570287322998047,
+      "grad_norm": 0.8075382113456726,
+      "learning_rate": 5.9768361406491533e-05,
+      "loss": 0.3967801284790039,
       "step": 22800
     },
     {
       "epoch": 2.1406666666666667,
-      "grad_norm": 4.561542510986328,
-      "learning_rate": 1.4551250397078253e-05,
-      "loss": 1.5680990600585938,
+      "grad_norm": 0.7892723679542542,
+      "learning_rate": 5.8205001588313014e-05,
+      "loss": 0.39166572570800784,
       "step": 22900
     },
     {
       "epoch": 2.144,
-      "grad_norm": 4.693943500518799,
-      "learning_rate": 1.4164717756691176e-05,
-      "loss": 1.5905775451660156,
+      "grad_norm": 0.7251806855201721,
+      "learning_rate": 5.6658871026764704e-05,
+      "loss": 0.3836771011352539,
       "step": 23000
     },
     {
       "epoch": 2.1473333333333335,
-      "grad_norm": 5.039695739746094,
-      "learning_rate": 1.3782539397351418e-05,
-      "loss": 1.5886349487304687,
+      "grad_norm": 0.782537579536438,
+      "learning_rate": 5.513015758940567e-05,
+      "loss": 0.39198822021484375,
       "step": 23100
     },
     {
       "epoch": 2.1506666666666665,
-      "grad_norm": 5.3282790184021,
-      "learning_rate": 1.340476175686784e-05,
-      "loss": 1.5803268432617188,
+      "grad_norm": 0.8268524408340454,
+      "learning_rate": 5.361904702747136e-05,
+      "loss": 0.38895458221435547,
       "step": 23200
     },
     {
       "epoch": 2.154,
-      "grad_norm": 4.672342777252197,
-      "learning_rate": 1.3031430738325778e-05,
-      "loss": 1.5878211975097656,
+      "grad_norm": 0.7508871555328369,
+      "learning_rate": 5.212572295330311e-05,
+      "loss": 0.40076961517333987,
       "step": 23300
     },
     {
       "epoch": 2.1573333333333333,
-      "grad_norm": 4.686863899230957,
-      "learning_rate": 1.2662591704509547e-05,
-      "loss": 1.611242218017578,
+      "grad_norm": 0.6501144170761108,
+      "learning_rate": 5.065036681803819e-05,
+      "loss": 0.4024519348144531,
       "step": 23400
     },
     {
       "epoch": 2.1606666666666667,
-      "grad_norm": 5.2701568603515625,
-      "learning_rate": 1.2298289472390417e-05,
-      "loss": 1.5942909240722656,
+      "grad_norm": 0.6458573937416077,
+      "learning_rate": 4.919315788956167e-05,
+      "loss": 0.4000421142578125,
       "step": 23500
     },
     {
       "epoch": 2.164,
-      "grad_norm": 4.708863258361816,
-      "learning_rate": 1.1938568307680964e-05,
-      "loss": 1.6314300537109374,
+      "grad_norm": 0.732280433177948,
+      "learning_rate": 4.775427323072386e-05,
+      "loss": 0.39758953094482424,
       "step": 23600
     },
     {
       "epoch": 2.1673333333333336,
-      "grad_norm": 4.914238929748535,
-      "learning_rate": 1.1583471919456506e-05,
-      "loss": 1.6009869384765625,
+      "grad_norm": 0.6772671937942505,
+      "learning_rate": 4.6333887677826025e-05,
+      "loss": 0.39006378173828127,
       "step": 23700
     },
     {
       "epoch": 2.1706666666666665,
-      "grad_norm": 5.067790508270264,
-      "learning_rate": 1.1233043454844016e-05,
-      "loss": 1.5916122436523437,
+      "grad_norm": 0.7737159729003906,
+      "learning_rate": 4.4932173819376064e-05,
+      "loss": 0.3929125213623047,
       "step": 23800
     },
     {
       "epoch": 2.174,
-      "grad_norm": 5.132817268371582,
-      "learning_rate": 1.0887325493779405e-05,
-      "loss": 1.6167672729492188,
+      "grad_norm": 0.8250499963760376,
+      "learning_rate": 4.354930197511762e-05,
+      "loss": 0.396903076171875,
       "step": 23900
     },
     {
       "epoch": 2.1773333333333333,
-      "grad_norm": 5.293063640594482,
-      "learning_rate": 1.0546360043833747e-05,
-      "loss": 1.6026080322265626,
+      "grad_norm": 0.7866168022155762,
+      "learning_rate": 4.218544017533499e-05,
+      "loss": 0.387851676940918,
       "step": 24000
     },
     {
       "epoch": 2.1806666666666668,
-      "grad_norm": 5.52355432510376,
-      "learning_rate": 1.0210188535108995e-05,
-      "loss": 1.6140333557128905,
+      "grad_norm": 0.7040881514549255,
+      "learning_rate": 4.084075414043598e-05,
+      "loss": 0.3934907913208008,
       "step": 24100
     },
     {
       "epoch": 2.184,
-      "grad_norm": 5.356894016265869,
-      "learning_rate": 9.878851815203882e-06,
-      "loss": 1.6005001831054688,
+      "grad_norm": 0.7672249674797058,
+      "learning_rate": 3.951540726081553e-05,
+      "loss": 0.3852077484130859,
       "step": 24200
     },
     {
       "epoch": 2.187333333333333,
-      "grad_norm": 5.4585771560668945,
-      "learning_rate": 9.55239014425055e-06,
-      "loss": 1.6022555541992187,
+      "grad_norm": 0.7783628702163696,
+      "learning_rate": 3.82095605770022e-05,
+      "loss": 0.3871307373046875,
       "step": 24300
     },
     {
       "epoch": 2.1906666666666665,
-      "grad_norm": 5.152921676635742,
-      "learning_rate": 9.230843190022725e-06,
-      "loss": 1.6190740966796875,
+      "grad_norm": 0.8870095610618591,
+      "learning_rate": 3.69233727600909e-05,
+      "loss": 0.3772114181518555,
       "step": 24400
     },
     {
       "epoch": 2.194,
-      "grad_norm": 5.397445201873779,
-      "learning_rate": 8.914250023115672e-06,
-      "loss": 1.616868438720703,
+      "grad_norm": 0.738532304763794,
+      "learning_rate": 3.565700009246269e-05,
+      "loss": 0.3705712890625,
       "step": 24500
     },
     {
       "epoch": 2.1973333333333334,
-      "grad_norm": 5.553057670593262,
-      "learning_rate": 8.602649112198796e-06,
-      "loss": 1.621106414794922,
+      "grad_norm": 0.6864748001098633,
+      "learning_rate": 3.4410596448795184e-05,
+      "loss": 0.383516845703125,
       "step": 24600
     },
     {
       "epoch": 2.2006666666666668,
-      "grad_norm": 4.983733177185059,
-      "learning_rate": 8.296078319341443e-06,
-      "loss": 1.6104743957519532,
+      "grad_norm": 0.6553505063056946,
+      "learning_rate": 3.318431327736577e-05,
+      "loss": 0.38183685302734377,
       "step": 24700
     },
     {
       "epoch": 2.204,
-      "grad_norm": 4.88228702545166,
-      "learning_rate": 7.994574895412293e-06,
-      "loss": 1.6215898132324218,
+      "grad_norm": 0.6033626794815063,
+      "learning_rate": 3.197829958164917e-05,
+      "loss": 0.3845671081542969,
       "step": 24800
     },
     {
       "epoch": 2.207333333333333,
-      "grad_norm": 5.843441963195801,
-      "learning_rate": 7.698175475553077e-06,
-      "loss": 1.613902587890625,
+      "grad_norm": 0.7695163488388062,
+      "learning_rate": 3.079270190221231e-05,
+      "loss": 0.3796828079223633,
       "step": 24900
     },
     {
       "epoch": 2.2106666666666666,
-      "grad_norm": 5.504673957824707,
-      "learning_rate": 7.40691607472715e-06,
-      "loss": 1.623016357421875,
+      "grad_norm": 0.7200196385383606,
+      "learning_rate": 2.96276642989086e-05,
+      "loss": 0.38924953460693357,
       "step": 25000
     },
     {
       "epoch": 2.2106666666666666,
-      "eval_loss": 1.500394344329834,
-      "eval_masked_accuracy": 0.6250727166957534,
-      "eval_runtime": 6.4092,
-      "eval_samples_per_second": 639.084,
-      "eval_steps_per_second": 2.496,
+      "eval_loss": 0.3823796212673187,
+      "eval_masked_accuracy": 0.878150720164609,
+      "eval_runtime": 3.0159,
+      "eval_samples_per_second": 1358.12,
+      "eval_steps_per_second": 5.305,
       "step": 25000
     },
     {
       "epoch": 2.214,
-      "grad_norm": 5.493242263793945,
-      "learning_rate": 7.120832083343337e-06,
-      "loss": 1.6365603637695312,
+      "grad_norm": 0.7287200689315796,
+      "learning_rate": 2.848332833337335e-05,
+      "loss": 0.3876392364501953,
       "step": 25100
     },
     {
       "epoch": 2.2173333333333334,
-      "grad_norm": 5.239762783050537,
-      "learning_rate": 6.839958262955709e-06,
-      "loss": 1.6364859008789063,
+      "grad_norm": 0.7604002952575684,
+      "learning_rate": 2.7359833051822837e-05,
+      "loss": 0.3845469284057617,
       "step": 25200
     },
     {
       "epoch": 2.220666666666667,
-      "grad_norm": 5.058232307434082,
-      "learning_rate": 6.564328742039783e-06,
-      "loss": 1.641663818359375,
+      "grad_norm": 0.6634399890899658,
+      "learning_rate": 2.625731496815913e-05,
+      "loss": 0.3846179962158203,
       "step": 25300
     },
     {
       "epoch": 2.224,
-      "grad_norm": 5.345419406890869,
-      "learning_rate": 6.293977011845648e-06,
-      "loss": 1.6362882995605468,
+      "grad_norm": 0.7272804975509644,
+      "learning_rate": 2.5175908047382592e-05,
+      "loss": 0.39008934020996094,
       "step": 25400
     },
     {
       "epoch": 2.227333333333333,
-      "grad_norm": 5.698750019073486,
-      "learning_rate": 6.0289359223284434e-06,
-      "loss": 1.6450096130371095,
+      "grad_norm": 0.714460015296936,
+      "learning_rate": 2.4115743689313774e-05,
+      "loss": 0.3927734756469727,
       "step": 25500
     },
     {
       "epoch": 2.2306666666666666,
-      "grad_norm": 5.43581485748291,
-      "learning_rate": 5.769237678156897e-06,
-      "loss": 1.6310890197753907,
+      "grad_norm": 0.8357155919075012,
+      "learning_rate": 2.3076950712627588e-05,
+      "loss": 0.38554725646972654,
       "step": 25600
     },
     {
       "epoch": 2.234,
-      "grad_norm": 5.170186996459961,
-      "learning_rate": 5.514913834800134e-06,
-      "loss": 1.6681497192382813,
+      "grad_norm": 0.7323917150497437,
+      "learning_rate": 2.2059655339200536e-05,
+      "loss": 0.3921313095092773,
       "step": 25700
     },
     {
       "epoch": 2.2373333333333334,
-      "grad_norm": 5.491499423980713,
-      "learning_rate": 5.265995294693471e-06,
-      "loss": 1.6503074645996094,
+      "grad_norm": 0.7257605791091919,
+      "learning_rate": 2.1063981178773884e-05,
+      "loss": 0.38508140563964843,
       "step": 25800
     },
     {
       "epoch": 2.240666666666667,
-      "grad_norm": 4.995376110076904,
-      "learning_rate": 5.022512303483451e-06,
-      "loss": 1.6396957397460938,
+      "grad_norm": 0.7842901945114136,
+      "learning_rate": 2.0090049213933803e-05,
+      "loss": 0.39042671203613283,
       "step": 25900
     },
     {
       "epoch": 2.2439999999999998,
-      "grad_norm": 5.047297477722168,
-      "learning_rate": 4.784494446352833e-06,
-      "loss": 1.658381805419922,
+      "grad_norm": 0.6478227376937866,
+      "learning_rate": 1.9137977785411333e-05,
+      "loss": 0.3918688201904297,
       "step": 26000
     },
     {
       "epoch": 2.247333333333333,
-      "grad_norm": 5.690036296844482,
-      "learning_rate": 4.551970644425707e-06,
-      "loss": 1.6450552368164062,
+      "grad_norm": 0.7710049748420715,
+      "learning_rate": 1.820788257770283e-05,
+      "loss": 0.3908076858520508,
       "step": 26100
     },
     {
       "epoch": 2.2506666666666666,
-      "grad_norm": 5.233314514160156,
-      "learning_rate": 4.324969151253317e-06,
-      "loss": 1.6168133544921874,
+      "grad_norm": 0.7351255416870117,
+      "learning_rate": 1.7299876605013267e-05,
+      "loss": 0.37932159423828127,
       "step": 26200
     },
     {
       "epoch": 2.254,
-      "grad_norm": 5.15321159362793,
-      "learning_rate": 4.10351754938107e-06,
-      "loss": 1.615177764892578,
+      "grad_norm": 0.6514462232589722,
+      "learning_rate": 1.641407019752428e-05,
+      "loss": 0.38645931243896486,
       "step": 26300
     },
     {
       "epoch": 2.2573333333333334,
-      "grad_norm": 6.050407886505127,
-      "learning_rate": 3.887642746997017e-06,
-      "loss": 1.6591305541992187,
+      "grad_norm": 0.7680016160011292,
+      "learning_rate": 1.5550570987988067e-05,
+      "loss": 0.381392707824707,
       "step": 26400
     },
     {
       "epoch": 2.260666666666667,
-      "grad_norm": 5.440352439880371,
-      "learning_rate": 3.6773709746622363e-06,
-      "loss": 1.6324659729003905,
+      "grad_norm": 0.7839369177818298,
+      "learning_rate": 1.4709483898648945e-05,
+      "loss": 0.37340118408203127,
       "step": 26500
     },
     {
       "epoch": 2.2640000000000002,
-      "grad_norm": 5.16418981552124,
-      "learning_rate": 3.472727782123697e-06,
-      "loss": 1.634721221923828,
+      "grad_norm": 0.6791050434112549,
+      "learning_rate": 1.3890911128494788e-05,
+      "loss": 0.38508747100830076,
       "step": 26600
     },
     {
       "epoch": 2.267333333333333,
-      "grad_norm": 5.912958145141602,
-      "learning_rate": 3.273738035209678e-06,
-      "loss": 1.6322593688964844,
+      "grad_norm": 0.6227408051490784,
+      "learning_rate": 1.3094952140838712e-05,
+      "loss": 0.3770198059082031,
       "step": 26700
     },
     {
       "epoch": 2.2706666666666666,
-      "grad_norm": 6.0202155113220215,
-      "learning_rate": 3.0804259128083936e-06,
-      "loss": 1.6233775329589843,
+      "grad_norm": 0.6502448916435242,
+      "learning_rate": 1.2321703651233574e-05,
+      "loss": 0.37348697662353514,
       "step": 26800
     },
     {
       "epoch": 2.274,
-      "grad_norm": 5.666730880737305,
-      "learning_rate": 2.8928149039300524e-06,
-      "loss": 1.6298930358886718,
+      "grad_norm": 0.6978487968444824,
+      "learning_rate": 1.157125961572021e-05,
+      "loss": 0.3783902359008789,
       "step": 26900
     },
     {
       "epoch": 2.2773333333333334,
-      "grad_norm": 5.367899417877197,
-      "learning_rate": 2.7109278048527752e-06,
-      "loss": 1.6212306213378906,
+      "grad_norm": 0.7682250738143921,
+      "learning_rate": 1.0843711219411101e-05,
+      "loss": 0.3767197799682617,
       "step": 27000
     },
     {
       "epoch": 2.280666666666667,
-      "grad_norm": 5.847186088562012,
-      "learning_rate": 2.5347867163526384e-06,
-      "loss": 1.6395831298828125,
+      "grad_norm": 0.6915271878242493,
+      "learning_rate": 1.0139146865410554e-05,
+      "loss": 0.3798466110229492,
       "step": 27100
     },
     {
       "epoch": 2.284,
-      "grad_norm": 5.451877117156982,
-      "learning_rate": 2.364413041018232e-06,
-      "loss": 1.6114913940429687,
+      "grad_norm": 0.819548487663269,
+      "learning_rate": 9.457652164072928e-06,
+      "loss": 0.3809815979003906,
       "step": 27200
     },
     {
       "epoch": 2.287333333333333,
-      "grad_norm": 5.2447190284729,
-      "learning_rate": 2.1998274806501074e-06,
-      "loss": 1.6211642456054687,
+      "grad_norm": 0.7806178331375122,
+      "learning_rate": 8.79930992260043e-06,
+      "loss": 0.3813274383544922,
       "step": 27300
     },
     {
       "epoch": 2.2906666666666666,
-      "grad_norm": 5.420331001281738,
-      "learning_rate": 2.0410500337453176e-06,
-      "loss": 1.5998692321777344,
+      "grad_norm": 0.9139585494995117,
+      "learning_rate": 8.16420013498127e-06,
+      "loss": 0.38432167053222654,
       "step": 27400
     },
     {
       "epoch": 2.294,
-      "grad_norm": 5.493936538696289,
-      "learning_rate": 1.8880999930674215e-06,
-      "loss": 1.6333512878417968,
+      "grad_norm": 0.719249963760376,
+      "learning_rate": 7.552399972269686e-06,
+      "loss": 0.3763787078857422,
       "step": 27500
     },
     {
       "epoch": 2.2973333333333334,
-      "grad_norm": 5.61967658996582,
-      "learning_rate": 1.7409959433022871e-06,
-      "loss": 1.637555389404297,
+      "grad_norm": 0.7366885542869568,
+      "learning_rate": 6.963983773209149e-06,
+      "loss": 0.3752469253540039,
       "step": 27600
     },
     {
       "epoch": 2.300666666666667,
-      "grad_norm": 5.362664699554443,
-      "learning_rate": 1.5997557587998868e-06,
-      "loss": 1.597519073486328,
+      "grad_norm": 0.7235630750656128,
+      "learning_rate": 6.399023035199547e-06,
+      "loss": 0.37348270416259766,
       "step": 27700
     },
     {
       "epoch": 2.304,
-      "grad_norm": 5.862389087677002,
-      "learning_rate": 1.4643966014023957e-06,
-      "loss": 1.6198776245117188,
+      "grad_norm": 0.6602711081504822,
+      "learning_rate": 5.857586405609583e-06,
+      "loss": 0.38024497985839845,
       "step": 27800
     },
     {
       "epoch": 2.3073333333333332,
-      "grad_norm": 5.315817356109619,
-      "learning_rate": 1.3349349183589155e-06,
-      "loss": 1.6214700317382813,
+      "grad_norm": 0.6910043954849243,
+      "learning_rate": 5.339739673435662e-06,
+      "loss": 0.3761766815185547,
       "step": 27900
     },
     {
       "epoch": 2.3106666666666666,
-      "grad_norm": 5.45588493347168,
-      "learning_rate": 1.2113864403269836e-06,
-      "loss": 1.6254971313476563,
+      "grad_norm": 0.7024565935134888,
+      "learning_rate": 4.8455457613079345e-06,
+      "loss": 0.37851264953613284,
       "step": 28000
     },
     {
       "epoch": 3.002666666666667,
-      "grad_norm": 5.6935906410217285,
-      "learning_rate": 1.0937661794611864e-06,
-      "loss": 1.4845796203613282,
+      "grad_norm": 0.9012552499771118,
+      "learning_rate": 4.375064717844746e-06,
+      "loss": 0.3574191665649414,
       "step": 28100
     },
     {
       "epoch": 3.006,
-      "grad_norm": 5.128046989440918,
-      "learning_rate": 9.820884275890286e-07,
-      "loss": 1.456612548828125,
+      "grad_norm": 0.680111825466156,
+      "learning_rate": 3.928353710356114e-06,
+      "loss": 0.3549924087524414,
       "step": 28200
     },
     {
       "epoch": 3.009333333333333,
-      "grad_norm": 5.23906946182251,
-      "learning_rate": 8.763667544743836e-07,
-      "loss": 1.4703726196289062,
+      "grad_norm": 0.79486083984375,
+      "learning_rate": 3.5054670178975346e-06,
+      "loss": 0.35243122100830077,
       "step": 28300
     },
     {
       "epoch": 3.0126666666666666,
-      "grad_norm": 5.33555793762207,
-      "learning_rate": 7.766140061686522e-07,
-      "loss": 1.4658856201171875,
+      "grad_norm": 0.6971395015716553,
+      "learning_rate": 3.1064560246746088e-06,
+      "loss": 0.34856960296630857,
       "step": 28400
     },
     {
       "epoch": 3.016,
-      "grad_norm": 5.468862533569336,
-      "learning_rate": 6.828423034498488e-07,
-      "loss": 1.4503196716308593,
+      "grad_norm": 0.7989268898963928,
+      "learning_rate": 2.7313692137993952e-06,
+      "loss": 0.35404052734375,
       "step": 28500
     },
     {
       "epoch": 3.0193333333333334,
-      "grad_norm": 5.648383140563965,
-      "learning_rate": 5.950630403498469e-07,
-      "loss": 1.4786280822753906,
+      "grad_norm": 0.681338369846344,
+      "learning_rate": 2.3802521613993878e-06,
+      "loss": 0.3551434326171875,
       "step": 28600
     },
     {
       "epoch": 3.022666666666667,
-      "grad_norm": 5.5019659996032715,
-      "learning_rate": 5.132868827698978e-07,
-      "loss": 1.4625227355957031,
+      "grad_norm": 0.6259098052978516,
+      "learning_rate": 2.053147531079591e-06,
+      "loss": 0.3543955230712891,
       "step": 28700
     },
     {
       "epoch": 3.026,
-      "grad_norm": 5.69915771484375,
-      "learning_rate": 4.375237671846333e-07,
-      "loss": 1.4489666748046874,
+      "grad_norm": 0.6789717078208923,
+      "learning_rate": 1.7500950687385332e-06,
+      "loss": 0.3479387664794922,
       "step": 28800
     },
     {
       "epoch": 3.029333333333333,
-      "grad_norm": 5.224785327911377,
-      "learning_rate": 3.677828994347154e-07,
-      "loss": 1.4654063415527343,
+      "grad_norm": 0.6986877918243408,
+      "learning_rate": 1.4711315977388617e-06,
+      "loss": 0.35285072326660155,
       "step": 28900
     },
     {
       "epoch": 3.0326666666666666,
-      "grad_norm": 5.944694519042969,
-      "learning_rate": 3.040727536082366e-07,
-      "loss": 1.4619117736816407,
+      "grad_norm": 0.8804546594619751,
+      "learning_rate": 1.2162910144329464e-06,
+      "loss": 0.3571059799194336,
       "step": 29000
     },
     {
       "epoch": 3.036,
-      "grad_norm": 5.931241989135742,
-      "learning_rate": 2.464010710110598e-07,
-      "loss": 1.464918975830078,
+      "grad_norm": 0.7633663415908813,
+      "learning_rate": 9.856042840442393e-07,
+      "loss": 0.3522502517700195,
       "step": 29100
     },
     {
       "epoch": 3.0393333333333334,
-      "grad_norm": 5.546744346618652,
-      "learning_rate": 1.9477485922618222e-07,
-      "loss": 1.4493109130859374,
+      "grad_norm": 0.8087472319602966,
+      "learning_rate": 7.790994369047289e-07,
+      "loss": 0.3503116607666016,
       "step": 29200
     },
     {
       "epoch": 3.042666666666667,
-      "grad_norm": 5.485001564025879,
-      "learning_rate": 1.4920039126225527e-07,
-      "loss": 1.4527821350097656,
+      "grad_norm": 0.6872431039810181,
+      "learning_rate": 5.968015650490211e-07,
+      "loss": 0.34913509368896484,
       "step": 29300
     },
     {
       "epoch": 3.046,
-      "grad_norm": 5.745002746582031,
-      "learning_rate": 1.0968320479136095e-07,
-      "loss": 1.46337646484375,
+      "grad_norm": 0.7077915668487549,
+      "learning_rate": 4.387328191654438e-07,
+      "loss": 0.34669517517089843,
       "step": 29400
     },
     {
       "epoch": 3.0493333333333332,
-      "grad_norm": 4.96608829498291,
-      "learning_rate": 7.622810147614456e-08,
-      "loss": 1.4616835021972656,
+      "grad_norm": 0.721786379814148,
+      "learning_rate": 3.0491240590457823e-07,
+      "loss": 0.3519774627685547,
       "step": 29500
     },
     {
       "epoch": 3.0526666666666666,
-      "grad_norm": 4.910149097442627,
-      "learning_rate": 4.883914638636467e-08,
-      "loss": 1.4573419189453125,
+      "grad_norm": 0.7876111268997192,
+      "learning_rate": 1.9535658554545867e-07,
+      "loss": 0.34955368041992185,
       "step": 29600
     },
     {
       "epoch": 3.056,
-      "grad_norm": 5.22091007232666,
-      "learning_rate": 2.7519667504971593e-08,
-      "loss": 1.4733062744140626,
+      "grad_norm": 0.5993502140045166,
+      "learning_rate": 1.1007867001988637e-07,
+      "loss": 0.3491347885131836,
       "step": 29700
     },
     {
       "epoch": 3.0593333333333335,
-      "grad_norm": 5.838210582733154,
-      "learning_rate": 1.2272255323708681e-08,
-      "loss": 1.49260498046875,
+      "grad_norm": 0.7406614422798157,
+      "learning_rate": 4.9089021294834726e-08,
+      "loss": 0.35303794860839843,
       "step": 29800
     },
     {
       "epoch": 3.062666666666667,
-      "grad_norm": 5.322798252105713,
-      "learning_rate": 3.0987625283473987e-09,
-      "loss": 1.4649183654785156,
+      "grad_norm": 0.7397328615188599,
+      "learning_rate": 1.2395050113389595e-08,
+      "loss": 0.34634124755859375,
       "step": 29900
     },
     {
       "epoch": 3.066,
-      "grad_norm": 6.612534999847412,
-      "learning_rate": 3.0377360626943076e-13,
-      "loss": 1.4366880798339843,
+      "grad_norm": 0.7702119946479797,
+      "learning_rate": 1.215094425077723e-12,
+      "loss": 0.34748374938964843,
       "step": 30000
     },
     {
       "epoch": 3.066,
-      "eval_loss": 1.5919160842895508,
-      "eval_masked_accuracy": 0.6138884886903905,
-      "eval_runtime": 6.0117,
-      "eval_samples_per_second": 681.335,
-      "eval_steps_per_second": 2.661,
+      "eval_loss": 0.37694016098976135,
+      "eval_masked_accuracy": 0.8775784753363228,
+      "eval_runtime": 2.9576,
+      "eval_samples_per_second": 1384.896,
+      "eval_steps_per_second": 5.41,
       "step": 30000
     },
     {
       "epoch": 3.066,
       "step": 30000,
-      "total_flos": 5.639177852861803e+16,
-      "train_loss": 2.222940887959798,
-      "train_runtime": 3365.6547,
-      "train_samples_per_second": 2281.874,
-      "train_steps_per_second": 8.914
+      "total_flos": 9.74439612930639e+16,
+      "train_loss": 0.6228721523284912,
+      "train_runtime": 2832.8619,
+      "train_samples_per_second": 2711.039,
+      "train_steps_per_second": 10.59
     },
     {
       "epoch": 3.066,
-      "eval_loss": 1.538759469985962,
-      "eval_masked_accuracy": 0.6228862973760932,
-      "eval_runtime": 5.9881,
-      "eval_samples_per_second": 684.026,
-      "eval_steps_per_second": 2.672,
+      "eval_loss": 0.3744058609008789,
+      "eval_masked_accuracy": 0.8805643239064613,
+      "eval_runtime": 2.8656,
+      "eval_samples_per_second": 1429.379,
+      "eval_steps_per_second": 5.584,
       "step": 30000
     }
   ],
@@ -2199,7 +2199,7 @@
       "attributes": {}
     }
   },
-  "total_flos": 5.639177852861803e+16,
+  "total_flos": 9.74439612930639e+16,
   "train_batch_size": 256,
   "trial_name": null,
   "trial_params": null