Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+text_to_code.json filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,232 @@
+---
+license: apache-2.0
+
+base_model:
+  - meta-llama/Meta-Llama-3-8B-Instruct
+
+language:
+  - es
+
+tags:
+  - BEL
+  - retrieval
+  - entity-retrieval
+  - named-entity-disambiguation
+  - entity-disambiguation
+  - named-entity-linking
+  - entity-linking
+  - text2text-generation
+  - biomedical
+  - healthcare
+  - synthetic-data
+  - causal-lm
+  - llm
+
+library_name: transformers
+finetuning_task:
+  - text2text-generation
+  - entity-linking
+---
+
+
+# SynCABEL: Synthetic Contextualized Augmentation for Biomedical Entity Linking
+
+## SynCABEL
+
+**SynCABEL** is a novel framework that addresses data scarcity in biomedical entity linking through **synthetic data generation**. The method, introduced in our [paper]
+
+## SynCABEL (SPACCC Edition)
+
+This is a **finetuned version of LLaMA-3-8B** trained on **SPACCC** using **[SynthSPACCC](https://huggingface.co/datasets/AnonymousARR42/SynCABEL)** (our synthetic dataset generated via the SynCABEL framework).
+
+| |  |
+|--------|---------|
+| **Base Model** | [meta-llama/Meta-Llama-3-8B-Instruct](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct) |
+| **Training Data** | SPACCC (real) + [SynthSPACCC](https://huggingface.co/datasets/AnonymousARR42/SynCABEL) (synthetic) |
+| **Fine-tuning** | [Supervised Fine-Tuning](https://huggingface.co/docs/trl/en/sft_trainer) |
+
+## Training Data Composition
+
+The model is trained on a mix of **human-annotated** and **synthetic** data:
+
+```
+SPACCC (human)   :  27,799 examples
+SynthSPACCC (synthetic)  : 1,813,463 examples
+```
+
+To ensure balanced learning, **human data is upsampled during training** so that each batch contains:
+
+```
+50% human-annotated data
+50% synthetic data
+```
+
+In other words, although SynthMM is larger, the model always sees a **1:1 ratio of human to synthetic examples**, preventing synthetic data from overwhelming human supervision.
+
+
+## Usage
+
+
+### Loading
+```python
+import torch
+from transformers import AutoModelForCausalLM
+
+# Load the model (requires trust_remote_code for custom architecture)
+model = AutoModelForCausalLM.from_pretrained(
+    "AnonymousARR42/SynCABEL_SPACCC",
+    trust_remote_code=True,
+    device_map="auto"
+)
+```
+
+### Unconstrained Generation
+```python
+# Let the model freely generate concept names
+sentences = [
+    "El paciente con [embolia pulmonar masiva]{ENFERMEDAD} presentó signos de dificultad respiratoria.",
+    "El paciente se sometió a una [angioplastia coronaria]{PROCEDIMIENTO} para restaurar el flujo sanguíneo."
+]
+
+results = model.sample(
+    sentences=sentences,
+    constrained=False,
+    num_beams=2,
+)
+
+for i, beam_results in enumerate(results):
+    print(f"Input: {sentences[i]}")
+
+    mention = beam_results[0]["mention"]
+    print(f"Mention: {mention}")
+
+    for j, result in enumerate(beam_results):
+        print(
+            f"Beam {j+1}:\n"
+            f"Predicted concept name:{result['pred_concept_name']}\n"
+            f"Predicted code: {result['pred_concept_code']}\n"
+            f"Beam score: {result['beam_score']:.3f}\n"
+        )
+```
+
+**Output:**
+```
+Input: El paciente con [embolia pulmonar masiva]{ENFERMEDAD} presentó signos de dificultad respiratoria.
+Mention: embolia pulmonar masiva
+Beam 1:
+Predicted concept name:tromboembolia pulmonar masiva aguda
+Predicted code: NO_CODE
+Beam score: 0.818
+
+Beam 2:
+Predicted concept name:tromboembolia masiva
+Predicted code: 58417008
+Beam score: 0.816
+
+Input: El paciente se sometió a una [angioplastia coronaria]{PROCEDIMIENTO} para restaurar el flujo sanguíneo.
+Mention: angioplastia coronaria
+Beam 1:
+Predicted concept name:operaciones transluminales en arteria coronaria
+Predicted code: NO_CODE
+Beam score: 0.764
+
+Beam 2:
+Predicted concept name:procedimiento en arteria coronaria
+Predicted code: NO_CODE
+Beam score: 0.728
+```
+
+### Constrained Decoding (Recommended for Entity Linking)
+```python
+# Constrained to valid biomedical concepts
+sentences = [
+    "El paciente con [embolia pulmonar masiva]{ENFERMEDAD} presentó signos de dificultad respiratoria.",
+    "El paciente se sometió a una [angioplastia coronaria]{PROCEDIMIENTO} para restaurar el flujo sanguíneo."
+]
+
+results = model.sample(
+    sentences=sentences,
+    constrained=True,
+    num_beams=2,
+)
+
+for i, beam_results in enumerate(results):
+    print(f"Input: {sentences[i]}")
+
+    mention = beam_results[0]["mention"]
+    print(f"Mention: {mention}")
+
+    for j, result in enumerate(beam_results):
+        print(
+            f"Beam {j+1}:\n"
+            f"Predicted concept name:{result['pred_concept_name']}\n"
+            f"Predicted code: {result['pred_concept_code']}\n"
+            f"Beam score: {result['beam_score']:.3f}\n"
+        )
+```
+
+**Output:**
+```
+Input: El paciente con [embolia pulmonar masiva]{ENFERMEDAD} presentó signos de dificultad respiratoria.
+Mention: embolia pulmonar masiva
+Beam 1:
+Predicted concept name:tromboembolia masiva
+Predicted code: 58417008
+Beam score: 0.816
+
+Beam 2:
+Predicted concept name:tromboembolia pulmonar aguda
+Predicted code: 707414004
+Beam score: 0.763
+
+Input: El paciente se sometió a una [angioplastia coronaria]{PROCEDIMIENTO} para restaurar el flujo sanguíneo.
+Mention: angioplastia coronaria
+Beam 1:
+Predicted concept name:operaciones transluminales en arteria pulmonar
+Predicted code: 175266007
+Beam score: 0.238
+
+Beam 2:
+Predicted concept name:operaciones transluminales en la arteria femoral o poplítea
+Predicted code: 265530008
+Beam score: 0.182
+```
+
+## Scores
+
+Entity linking performance (Recall@1) on biomedical benchmarks. The best results are shown in **bold**, the second-best results are <u>underlined</u>, and the "Average" column reports the mean score across the four benchmarks.
+
+| Model | MM-ST21PV<br>(english) | QUAERO-MEDLINE<br>(french) | QUAERO-EMEA<br>(french) | SPACCC<br>(spanish) | Avg. |
+| :--- | :---: | :---: | :---: | :---: | :---: |
+| SciSpacy | 53.8 | 40.5 | 37.1 | 13.2 | 36.2 |
+| SapBERT | 51.1 | 50.6 | 49.8 | 33.9 | 46.4 |
+| CODER-all | 56.6 | 58.7 | 58.1 | 43.7 | 54.3 |
+| SapBERT-all | 64.6 | 74.7 | 67.9 | 47.9 | 63.8 |
+| ArboEL | <u>74.5</u> | 70.9 | 62.8 | 49.0 | 64.2 |
+| mBART-large | 65.5 | 61.5 | 58.6 | 57.7 | 60.8 |
+| + Guided inference | 70.0 | 72.8 | 71.1 | 61.8 | 68.9 |
+| **+ SynCABEL (Our method)** | 71.5 | 77.1 | <u>75.3</u> | 64.0 | 72.0 |
+| Llama-3-8B | 69.0 | 66.4 | 65.5 | 59.9 | 65.2 |
+| + Guided inference | 74.4 | <u>77.5</u> | 72.9 | <u>64.2</u> | <u>72.3</u> |
+| **+ SynCABEL (Our method)** | **75.4** | **79.7** | **79.0** | **67.0** | **75.3** |
+
+Here, we provide the source repositories for the baselines:
+- [**SciSpacy**](https://github.com/allenai/scispacy)
+- [**SapBERT**](https://hf.co/cambridgeltl/SapBERT-from-PubMedBERT-fulltext)
+- [**SapBERT-all**](https://hf.co/cambridgeltl/SapBERT-UMLS-2020AB-all-lang-from-XLMR)
+- [**CODER-all**](https://hf.co/GanjinZero/coder_all)
+- [**ArboEL**](https://github.com/dhdhagar/arboEL)
+- [**mBART-large**](https://hf.co/facebook/mbart-large-50)
+- [**LLaMA-3-8B**](https://hf.co/meta-llama/Meta-Llama-3-8B-Instruct).
+
+
+### Speed and Memory
+
+| Model        | Model (GB) | Cand. (GB) | Speed (/s) |
+|--------------|------------|------------|------------|
+| SapBERT      | 2.1        | 20.1       | **575.5**  |
+| ArboEL       | **1.2**    | 7.1        | 38.9       |
+| mBART        | 2.3        | **5.4**    | 51.0       |
+| Llama-3-8B   | 28.6       | **5.4**    | 19.1       |
+
+*Measured on single H100 GPU, constrained decoding*

__init__.py

@@ -0,0 +1,4 @@
+# __init__.py
+from .syncabel import LLamaSynCABEL, LLamaSynCABELConfig
+
+__all__ = ["LLamaSynCABEL", "LLamaSynCABELConfig"]

candidate_trie.pkl

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

chat_template.jinja

@@ -0,0 +1,5 @@
+{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>
+
+'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>
+
+' }}

config.json

@@ -0,0 +1,35 @@
+{
+  "architectures": [
+    "LLamaSynCABEL"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 128000,
+  "dtype": "bfloat16",
+  "eos_token_id": 128001,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama_syncabel",
+  "auto_map": {
+    "AutoConfig": "syncabel.LLamaSynCABELConfig",
+    "AutoModelForCausalLM": "syncabel.LLamaSynCABEL"
+  },
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 8,
+  "pad_token_id": 128001,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 128257,
+  "lang": "es"
+}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 128000,
+  "eos_token_id": [
+    128001
+  ],
+  "pad_token_id": 128001,
+  "transformers_version": "4.57.1"
+}

guided_inference.py

@@ -0,0 +1,70 @@
+def _get_tgt_lang_token_id(tokenizer):
+    """Best-effort retrieval of target language token id from a HF tokenizer.
+
+    Returns None when not available.
+    """
+    # Some tokenizers (MBart, M2M, NLLB) expose `tgt_lang` and different ways to map to ids
+    tgt = getattr(tokenizer, "tgt_lang", None)
+    if not tgt:
+        return None
+
+    # Common mapping dict
+    try:
+        lang_code_to_id = getattr(tokenizer, "lang_code_to_id", None)
+        if isinstance(lang_code_to_id, dict) and tgt in lang_code_to_id:
+            return lang_code_to_id[tgt]
+    except Exception:
+        pass
+
+    return None
+
+
+def get_prefix_allowed_tokens_fn(
+    model,
+    sources: list[str],
+    prefix_templates: list[str],
+    sem_groups: list[str],
+    multiple_answers: bool = False,
+):
+    candidates_trie = model.candidate_trie  # type: ignore
+    sep_token_id = model.tokenizer.sep_token_id
+    eos_token_id = model.tokenizer.eos_token_id
+    pad_token_id = model.tokenizer.pad_token_id
+    tgt_lang_id = _get_tgt_lang_token_id(model.tokenizer)
+    prefix_templates = [model.tokenizer.encode(prefix) for prefix in prefix_templates]
+
+    def prefix_allowed_tokens_fn(batch_id, sent):
+        sent = sent.tolist()
+        prefix = prefix_templates[batch_id]
+        # Remove the prefix from the sent
+        index_sep = sent.index(sep_token_id)
+        sent = sent[index_sep + 1 :]
+        # Check if the prefix is present
+        prefix_len = len(prefix)
+        if sent[:prefix_len] == prefix:
+            sent = sent[prefix_len - 1 :]
+        else:
+            raise ValueError("Prefix not found in the generated sentence.")
+        if len(sent) > 1 and sent[-1] in [eos_token_id, pad_token_id]:
+            return [pad_token_id, eos_token_id]
+        sem_group = sem_groups[batch_id]
+        # Remove everything up to last sep_token_id and add prefix and tgt_lang_id
+        if multiple_answers and sep_token_id in sent:
+            sep_index = len(sent) - 1 - sent[::-1].index(sep_token_id)
+            if sep_index == len(sent) - 1:
+                # Start fresh with decoder start (and optional tgt language token)
+                sent = [prefix[-1]] + ([tgt_lang_id] if tgt_lang_id is not None else [])
+            else:
+                sent = (
+                    [prefix[-1]]
+                    + ([tgt_lang_id] if tgt_lang_id is not None else [])
+                    + sent[sep_index + 1 :]
+                )
+        trie_out = candidates_trie[
+            sem_group  # type: ignore
+        ].get(sent)
+        if multiple_answers and eos_token_id in trie_out:
+            trie_out = [sep_token_id] + trie_out
+        return trie_out
+
+    return prefix_allowed_tokens_fn

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,299 @@
+{
+  "metadata": {
+    "total_parameters": 8030269440,
+    "total_size": 16060538880
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00004-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.norm.weight": "model-00004-of-00004.safetensors"
+  }
+}

optimizer.pt

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

rng_state.pth

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

scheduler.pt

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

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|begin_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "<SEP>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

syncabel.py

@@ -0,0 +1,401 @@
+"""
+Core models for SynCABEL
+"""
+
+import json
+import logging
+import os
+import pickle
+import re
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from transformers import (
+    AutoTokenizer,
+    LlamaForCausalLM,
+    PretrainedConfig,
+)
+
+from .guided_inference import get_prefix_allowed_tokens_fn
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(levelname)s - %(message)s",
+)
+
+
+# Define a simple config class that inherits from PretrainedConfig
+class LLamaSynCABELConfig(PretrainedConfig):
+    model_type = "llama_syncabel"
+
+    def __init__(self, **kwargs):
+        # Ensure it has llama as base
+        kwargs.setdefault("model_type", "llama")
+        super().__init__(**kwargs)
+
+
+def chunk_it(seq, num):
+    assert num > 0
+    chunk_len = len(seq) // num
+    chunks = [seq[i * chunk_len : i * chunk_len + chunk_len] for i in range(num)]
+
+    diff = len(seq) - chunk_len * num
+    for i in range(diff):
+        chunks[i].append(seq[chunk_len * num + i])
+
+    return chunks
+
+
+def find_mention(text: str) -> str:
+    match = re.search(r"\[(.*?)\]", text)
+    if match:
+        return match.group(1).strip()
+    else:
+        raise ValueError("No mention found in the text.")
+
+
+def find_sem_group(text: str) -> str:
+    match = re.search(r"\{(.*?)\}", text)
+    if match:
+        return match.group(1).strip()
+    else:
+        raise ValueError("No group type found in the text.")
+
+
+def parse_prediction(
+    outputs: list[str],
+    sem_groups: list[str],
+    verb: str,
+    text_to_code: Optional[dict[str, dict[str, str]]] = None,
+    multiple_answers: bool = False,
+) -> tuple[list[str], list[str]]:
+    codes = []
+    predictions = []
+    for output, group in zip(outputs, sem_groups):
+        splits = output.split(f"] {verb}")  # type: ignore
+        if len(splits) > 1 and splits[1].strip():
+            prediction = splits[1].strip()
+            if text_to_code:
+                if multiple_answers:
+                    prediction_list = prediction.split("<SEP>")  # type: ignore
+                    code_list = []
+                    for pred in prediction_list:
+                        code_list.append(
+                            text_to_code[group].get(pred.strip(), "NO_CODE")
+                        )
+                    code = "+".join(code_list)
+                else:
+                    code = text_to_code[group].get(prediction, "NO_CODE")
+            else:
+                code = "NO_CODE"
+        else:
+            print(
+                "IndexError: splitting failed or empty prediction, adding empty string as prediction."
+            )
+            print(f"Full text: {output}")  # type: ignore
+            prediction = "NO_PREDICTION"
+            code = "NO_CODE"
+        codes.append(code)
+        predictions.append(prediction)
+    return codes, predictions
+
+
+def compute_score(outputs, tokenizer, prefix_len=0):
+    sequences = outputs.sequences
+    scores = outputs.scores
+
+    N, total_len = sequences.shape
+    T = len(scores)
+
+    sequences = sequences[:, prefix_len : prefix_len + T]
+
+    if len(scores) > sequences.size(1):
+        scores = scores[: sequences.size(1)]
+
+    mask = (
+        (sequences != tokenizer.pad_token_id)
+        & (sequences != tokenizer.eos_token_id)
+        & (sequences != tokenizer.bos_token_id)
+    )
+
+    logprob_steps = []
+    for t, logits in enumerate(scores):
+        log_probs_t = F.log_softmax(logits, dim=-1)
+        token_t = sequences[:, t]
+        idx = torch.arange(N)
+        logprob_steps.append(log_probs_t[idx, token_t])
+
+    logprobs = torch.stack(logprob_steps, dim=1)
+    logprobs.masked_fill_(~mask, 0)
+
+    lengths = mask.sum(dim=1).clamp(min=1)
+    confidence = torch.exp(logprobs.sum(dim=1) / lengths)
+
+    return confidence.tolist()
+
+
+def skip_undesired_tokens(outputs, tokenizer):
+    sep_token = tokenizer.sep_token if tokenizer.sep_token is not None else None
+
+    if any("tag" in token for token in tokenizer.all_special_tokens):
+        tokens_to_remove = tokenizer.all_special_tokens[:-3]
+    elif any("{" in token for token in tokenizer.all_special_tokens):
+        tokens_to_remove = tokenizer.all_special_tokens[:-4]
+    else:
+        tokens_to_remove = tokenizer.all_special_tokens
+
+    if sep_token in tokens_to_remove:
+        tokens_to_remove = [tok for tok in tokens_to_remove if tok != sep_token]
+
+    cleaned_outputs = []
+    for sequence in outputs:
+        for token in tokens_to_remove:
+            sequence = sequence.replace(token, "")
+
+        if sep_token:
+            sequence = re.sub(rf"({re.escape(sep_token)})\s+", r"\1", sequence)
+
+        cleaned_outputs.append(sequence.strip())
+
+    return cleaned_outputs
+
+
+class LLamaSynCABEL(LlamaForCausalLM):
+    config_class = LLamaSynCABELConfig
+
+    def __init__(self, config, *args, **kwargs):
+        # Initialize the parent LlamaForCausalLM
+        super().__init__(config, *args, **kwargs)
+
+        # Store language from config
+        self.lang = getattr(config, "lang", "en")
+        self.text_to_code = None
+        self.candidate_trie = None
+        self.tokenizer = None
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        lang=None,
+        text_to_code_path=None,
+        candidate_trie_path=None,
+        **kwargs,
+    ):
+        # Remove custom kwargs before passing to parent
+        custom_kwargs = {
+            "lang": lang,
+            "text_to_code_path": text_to_code_path,
+            "candidate_trie_path": candidate_trie_path,
+        }
+
+        # Call parent's from_pretrained
+        model = super().from_pretrained(
+            pretrained_model_name_or_path,
+            *args,
+            **{k: v for k, v in kwargs.items() if k not in custom_kwargs},
+        )
+
+        # Set up tokenizer
+        model.tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path, use_fast=True
+        )
+        model.tokenizer.padding_side = "left"
+
+        # Set language: explicit override > config > default
+        if lang is not None:
+            model.lang = lang
+        elif hasattr(model.config, "lang"):
+            model.lang = model.config.lang
+        else:
+            model.lang = "en"
+
+        logger.info(f"Model language set to: {model.lang}")
+
+        # Load text_to_code
+        text_to_code_file_local = (
+            text_to_code_path
+            if text_to_code_path is not None
+            else os.path.join(pretrained_model_name_or_path, "text_to_code.json")
+        )
+        try:
+            if os.path.exists(text_to_code_file_local):
+                with open(text_to_code_file_local, encoding="utf-8") as f:
+                    model.text_to_code = json.load(f)
+                logger.info(
+                    f"Loaded text_to_code.json from local path: {text_to_code_file_local}"
+                )
+            else:
+                text_to_code_path_hf = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="text_to_code.json",
+                )
+                with open(text_to_code_path_hf, encoding="utf-8") as f:
+                    model.text_to_code = json.load(f)
+                logger.info(
+                    f"Loaded text_to_code.json from HF Hub: {text_to_code_path_hf}"
+                )
+        except Exception:
+            logger.warning("text_to_code.json not found (local or HF hub)")
+            model.text_to_code = None
+
+        # Load candidate_trie
+        candidate_trie_file_local = (
+            candidate_trie_path
+            if candidate_trie_path is not None
+            else os.path.join(pretrained_model_name_or_path, "candidate_trie.pkl")
+        )
+        try:
+            if os.path.exists(candidate_trie_file_local):
+                with open(candidate_trie_file_local, "rb") as f:
+                    model.candidate_trie = pickle.load(f)
+                logger.info(
+                    f"Loaded candidate_trie.pkl from local path: {candidate_trie_file_local}"
+                )
+            else:
+                candidate_trie_path_hf = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="candidate_trie.pkl",
+                )
+                with open(candidate_trie_path_hf, "rb") as f:
+                    model.candidate_trie = pickle.load(f)
+                logger.info(
+                    f"Loaded candidate_trie.pkl from HF Hub: {candidate_trie_path_hf}"
+                )
+        except Exception:
+            logger.warning("candidate_trie.pkl not found (local or HF hub)")
+            model.candidate_trie = None
+
+        return model
+
+    def sample(
+        self,
+        sentences: str | list[str],  # type: ignore
+        num_beams: int = 5,
+        constrained: bool = True,
+        multiple_answers: bool = False,
+        **kwargs,
+    ) -> list[list[dict[str, str]]]:
+
+        if isinstance(sentences, str):
+            sentences = [sentences]
+
+        if self.lang == "fr":
+            verb = "est"
+        elif self.lang == "en":
+            verb = "is"
+        elif self.lang == "es":
+            verb = "es"
+        else:
+            raise ValueError(f"Unsupported language: {self.lang}")
+
+        prefix_templates = []
+        complete_input_text = []
+        sem_groups = []
+        mentions = []
+        for sent in sentences:
+            sem_group = find_sem_group(sent)
+            mention = find_mention(sent)
+            prefix = f"[{mention}] {verb}"
+            complete_input = f"{sent}<SEP>{prefix}"
+            mentions.append(mention)
+            prefix_templates.append(prefix)
+            complete_input_text.append(complete_input)
+            sem_groups.append(sem_group)
+
+        input_args = {
+            k: v.to(self.device)
+            for k, v in self.tokenizer.batch_encode_plus(  # type: ignore
+                complete_input_text, padding="longest", return_tensors="pt"
+            ).items()
+        }
+
+        prefix_allowed_tokens_fn = None
+        if constrained:
+            if self.candidate_trie is None:
+                raise ValueError(
+                    "candidate_trie is not loaded in the model. Use constrained=False."
+                )
+            prefix_allowed_tokens_fn = get_prefix_allowed_tokens_fn(
+                self,
+                sentences,
+                prefix_templates,
+                sem_groups,
+                multiple_answers=multiple_answers,
+            )
+
+        outputs = self.generate(
+            **input_args,
+            max_new_tokens=128,
+            num_beams=num_beams,
+            num_return_sequences=num_beams,
+            output_scores=True,
+            return_dict_in_generate=True,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            **kwargs,
+        )
+
+        decoded_sequences = self.tokenizer.batch_decode(  # type: ignore
+            outputs.sequences,  # type: ignore
+            skip_special_tokens=False,
+            clean_up_tokenization_spaces=True,
+        )
+        cleaned_output_sequences = skip_undesired_tokens(
+            decoded_sequences,
+            self.tokenizer,
+        )
+
+        prefix_len = input_args["input_ids"].size(1)
+
+        sem_groups = [x for x in sem_groups for _ in range(num_beams)]
+        mentions = [x for x in mentions for _ in range(num_beams)]
+
+        codes, predictions = parse_prediction(
+            cleaned_output_sequences,
+            sem_groups,
+            verb,
+            self.text_to_code,
+            multiple_answers=multiple_answers,
+        )
+        scores = compute_score(outputs, self.tokenizer, prefix_len=prefix_len)
+        beam_scores = [
+            float(torch.exp(s)) if num_beams > 1 else float("nan")
+            for s in (
+                outputs.sequences_scores  # type: ignore
+                if num_beams > 1
+                else [torch.tensor(float("nan"))] * len(scores)
+            )
+        ]
+
+        outputs = chunk_it(
+            [
+                {
+                    "text": text,
+                    "mention": mention,
+                    "semantic_group": group,
+                    "pred_concept_name": prediction,
+                    "pred_concept_code": code,
+                    "score": score,
+                    "beam_score": beam_score,
+                }
+                for text, score, beam_score, code, prediction, mention, group in zip(
+                    cleaned_output_sequences,
+                    scores,
+                    beam_scores,
+                    codes,
+                    predictions,
+                    mentions,
+                    sem_groups,
+                )
+            ],
+            len(sentences),
+        )
+
+        return outputs
+
+    def encode(self, sentence):
+        return self.tokenizer.encode(sentence, return_tensors="pt")[0]  # type: ignore

text_to_code.json

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

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,2072 @@
+{
+  "added_tokens_decoder": {
+    "128000": {
+      "content": "<|begin_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128001": {
+      "content": "<|end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128002": {
+      "content": "<|reserved_special_token_0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128003": {
+      "content": "<|reserved_special_token_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128004": {
+      "content": "<|reserved_special_token_2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128005": {
+      "content": "<|reserved_special_token_3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128006": {
+      "content": "<|start_header_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128007": {
+      "content": "<|end_header_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128008": {
+      "content": "<|reserved_special_token_4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128009": {
+      "content": "<|eot_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128010": {
+      "content": "<|reserved_special_token_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128011": {
+      "content": "<|reserved_special_token_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128012": {
+      "content": "<|reserved_special_token_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128013": {
+      "content": "<|reserved_special_token_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128014": {
+      "content": "<|reserved_special_token_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128015": {
+      "content": "<|reserved_special_token_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128016": {
+      "content": "<|reserved_special_token_11|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128017": {
+      "content": "<|reserved_special_token_12|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128018": {
+      "content": "<|reserved_special_token_13|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128019": {
+      "content": "<|reserved_special_token_14|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128020": {
+      "content": "<|reserved_special_token_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128021": {
+      "content": "<|reserved_special_token_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128022": {
+      "content": "<|reserved_special_token_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128023": {
+      "content": "<|reserved_special_token_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128024": {
+      "content": "<|reserved_special_token_19|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128025": {
+      "content": "<|reserved_special_token_20|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128026": {
+      "content": "<|reserved_special_token_21|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128027": {
+      "content": "<|reserved_special_token_22|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128028": {
+      "content": "<|reserved_special_token_23|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128029": {
+      "content": "<|reserved_special_token_24|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128030": {
+      "content": "<|reserved_special_token_25|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128031": {
+      "content": "<|reserved_special_token_26|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128032": {
+      "content": "<|reserved_special_token_27|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128033": {
+      "content": "<|reserved_special_token_28|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128034": {
+      "content": "<|reserved_special_token_29|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128035": {
+      "content": "<|reserved_special_token_30|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128036": {
+      "content": "<|reserved_special_token_31|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128037": {
+      "content": "<|reserved_special_token_32|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128038": {
+      "content": "<|reserved_special_token_33|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128039": {
+      "content": "<|reserved_special_token_34|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128040": {
+      "content": "<|reserved_special_token_35|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128041": {
+      "content": "<|reserved_special_token_36|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128042": {
+      "content": "<|reserved_special_token_37|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128043": {
+      "content": "<|reserved_special_token_38|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128044": {
+      "content": "<|reserved_special_token_39|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128045": {
+      "content": "<|reserved_special_token_40|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128046": {
+      "content": "<|reserved_special_token_41|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128047": {
+      "content": "<|reserved_special_token_42|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128048": {
+      "content": "<|reserved_special_token_43|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128049": {
+      "content": "<|reserved_special_token_44|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128050": {
+      "content": "<|reserved_special_token_45|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128051": {
+      "content": "<|reserved_special_token_46|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128052": {
+      "content": "<|reserved_special_token_47|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128053": {
+      "content": "<|reserved_special_token_48|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128054": {
+      "content": "<|reserved_special_token_49|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128055": {
+      "content": "<|reserved_special_token_50|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128056": {
+      "content": "<|reserved_special_token_51|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128057": {
+      "content": "<|reserved_special_token_52|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128058": {
+      "content": "<|reserved_special_token_53|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128059": {
+      "content": "<|reserved_special_token_54|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128060": {
+      "content": "<|reserved_special_token_55|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128061": {
+      "content": "<|reserved_special_token_56|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128062": {
+      "content": "<|reserved_special_token_57|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128063": {
+      "content": "<|reserved_special_token_58|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128064": {
+      "content": "<|reserved_special_token_59|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128065": {
+      "content": "<|reserved_special_token_60|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128066": {
+      "content": "<|reserved_special_token_61|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128067": {
+      "content": "<|reserved_special_token_62|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128068": {
+      "content": "<|reserved_special_token_63|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128069": {
+      "content": "<|reserved_special_token_64|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128070": {
+      "content": "<|reserved_special_token_65|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128071": {
+      "content": "<|reserved_special_token_66|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128072": {
+      "content": "<|reserved_special_token_67|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128073": {
+      "content": "<|reserved_special_token_68|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128074": {
+      "content": "<|reserved_special_token_69|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128075": {
+      "content": "<|reserved_special_token_70|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128076": {
+      "content": "<|reserved_special_token_71|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128077": {
+      "content": "<|reserved_special_token_72|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128078": {
+      "content": "<|reserved_special_token_73|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128079": {
+      "content": "<|reserved_special_token_74|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128080": {
+      "content": "<|reserved_special_token_75|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128081": {
+      "content": "<|reserved_special_token_76|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128082": {
+      "content": "<|reserved_special_token_77|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128083": {
+      "content": "<|reserved_special_token_78|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128084": {
+      "content": "<|reserved_special_token_79|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128085": {
+      "content": "<|reserved_special_token_80|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128086": {
+      "content": "<|reserved_special_token_81|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128087": {
+      "content": "<|reserved_special_token_82|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128088": {
+      "content": "<|reserved_special_token_83|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128089": {
+      "content": "<|reserved_special_token_84|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128090": {
+      "content": "<|reserved_special_token_85|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128091": {
+      "content": "<|reserved_special_token_86|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128092": {
+      "content": "<|reserved_special_token_87|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128093": {
+      "content": "<|reserved_special_token_88|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128094": {
+      "content": "<|reserved_special_token_89|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128095": {
+      "content": "<|reserved_special_token_90|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128096": {
+      "content": "<|reserved_special_token_91|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128097": {
+      "content": "<|reserved_special_token_92|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128098": {
+      "content": "<|reserved_special_token_93|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128099": {
+      "content": "<|reserved_special_token_94|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128100": {
+      "content": "<|reserved_special_token_95|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128101": {
+      "content": "<|reserved_special_token_96|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128102": {
+      "content": "<|reserved_special_token_97|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128103": {
+      "content": "<|reserved_special_token_98|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128104": {
+      "content": "<|reserved_special_token_99|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128105": {
+      "content": "<|reserved_special_token_100|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128106": {
+      "content": "<|reserved_special_token_101|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128107": {
+      "content": "<|reserved_special_token_102|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128108": {
+      "content": "<|reserved_special_token_103|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128109": {
+      "content": "<|reserved_special_token_104|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128110": {
+      "content": "<|reserved_special_token_105|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128111": {
+      "content": "<|reserved_special_token_106|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128112": {
+      "content": "<|reserved_special_token_107|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128113": {
+      "content": "<|reserved_special_token_108|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128114": {
+      "content": "<|reserved_special_token_109|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128115": {
+      "content": "<|reserved_special_token_110|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128116": {
+      "content": "<|reserved_special_token_111|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128117": {
+      "content": "<|reserved_special_token_112|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128118": {
+      "content": "<|reserved_special_token_113|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128119": {
+      "content": "<|reserved_special_token_114|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128120": {
+      "content": "<|reserved_special_token_115|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128121": {
+      "content": "<|reserved_special_token_116|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128122": {
+      "content": "<|reserved_special_token_117|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128123": {
+      "content": "<|reserved_special_token_118|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128124": {
+      "content": "<|reserved_special_token_119|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128125": {
+      "content": "<|reserved_special_token_120|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128126": {
+      "content": "<|reserved_special_token_121|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128127": {
+      "content": "<|reserved_special_token_122|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128128": {
+      "content": "<|reserved_special_token_123|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128129": {
+      "content": "<|reserved_special_token_124|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128130": {
+      "content": "<|reserved_special_token_125|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128131": {
+      "content": "<|reserved_special_token_126|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128132": {
+      "content": "<|reserved_special_token_127|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128133": {
+      "content": "<|reserved_special_token_128|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128134": {
+      "content": "<|reserved_special_token_129|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128135": {
+      "content": "<|reserved_special_token_130|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128136": {
+      "content": "<|reserved_special_token_131|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128137": {
+      "content": "<|reserved_special_token_132|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128138": {
+      "content": "<|reserved_special_token_133|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128139": {
+      "content": "<|reserved_special_token_134|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128140": {
+      "content": "<|reserved_special_token_135|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128141": {
+      "content": "<|reserved_special_token_136|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128142": {
+      "content": "<|reserved_special_token_137|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128143": {
+      "content": "<|reserved_special_token_138|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128144": {
+      "content": "<|reserved_special_token_139|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128145": {
+      "content": "<|reserved_special_token_140|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128146": {
+      "content": "<|reserved_special_token_141|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128147": {
+      "content": "<|reserved_special_token_142|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128148": {
+      "content": "<|reserved_special_token_143|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128149": {
+      "content": "<|reserved_special_token_144|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128150": {
+      "content": "<|reserved_special_token_145|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128151": {
+      "content": "<|reserved_special_token_146|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128152": {
+      "content": "<|reserved_special_token_147|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128153": {
+      "content": "<|reserved_special_token_148|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128154": {
+      "content": "<|reserved_special_token_149|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128155": {
+      "content": "<|reserved_special_token_150|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128156": {
+      "content": "<|reserved_special_token_151|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128157": {
+      "content": "<|reserved_special_token_152|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128158": {
+      "content": "<|reserved_special_token_153|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128159": {
+      "content": "<|reserved_special_token_154|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128160": {
+      "content": "<|reserved_special_token_155|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128161": {
+      "content": "<|reserved_special_token_156|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128162": {
+      "content": "<|reserved_special_token_157|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128163": {
+      "content": "<|reserved_special_token_158|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128164": {
+      "content": "<|reserved_special_token_159|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128165": {
+      "content": "<|reserved_special_token_160|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128166": {
+      "content": "<|reserved_special_token_161|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128167": {
+      "content": "<|reserved_special_token_162|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128168": {
+      "content": "<|reserved_special_token_163|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128169": {
+      "content": "<|reserved_special_token_164|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128170": {
+      "content": "<|reserved_special_token_165|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128171": {
+      "content": "<|reserved_special_token_166|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128172": {
+      "content": "<|reserved_special_token_167|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128173": {
+      "content": "<|reserved_special_token_168|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128174": {
+      "content": "<|reserved_special_token_169|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128175": {
+      "content": "<|reserved_special_token_170|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128176": {
+      "content": "<|reserved_special_token_171|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128177": {
+      "content": "<|reserved_special_token_172|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128178": {
+      "content": "<|reserved_special_token_173|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128179": {
+      "content": "<|reserved_special_token_174|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128180": {
+      "content": "<|reserved_special_token_175|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128181": {
+      "content": "<|reserved_special_token_176|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128182": {
+      "content": "<|reserved_special_token_177|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128183": {
+      "content": "<|reserved_special_token_178|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128184": {
+      "content": "<|reserved_special_token_179|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128185": {
+      "content": "<|reserved_special_token_180|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128186": {
+      "content": "<|reserved_special_token_181|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128187": {
+      "content": "<|reserved_special_token_182|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128188": {
+      "content": "<|reserved_special_token_183|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128189": {
+      "content": "<|reserved_special_token_184|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128190": {
+      "content": "<|reserved_special_token_185|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128191": {
+      "content": "<|reserved_special_token_186|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128192": {
+      "content": "<|reserved_special_token_187|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128193": {
+      "content": "<|reserved_special_token_188|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128194": {
+      "content": "<|reserved_special_token_189|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128195": {
+      "content": "<|reserved_special_token_190|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128196": {
+      "content": "<|reserved_special_token_191|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128197": {
+      "content": "<|reserved_special_token_192|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128198": {
+      "content": "<|reserved_special_token_193|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128199": {
+      "content": "<|reserved_special_token_194|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128200": {
+      "content": "<|reserved_special_token_195|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128201": {
+      "content": "<|reserved_special_token_196|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128202": {
+      "content": "<|reserved_special_token_197|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128203": {
+      "content": "<|reserved_special_token_198|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128204": {
+      "content": "<|reserved_special_token_199|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128205": {
+      "content": "<|reserved_special_token_200|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128206": {
+      "content": "<|reserved_special_token_201|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128207": {
+      "content": "<|reserved_special_token_202|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128208": {
+      "content": "<|reserved_special_token_203|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128209": {
+      "content": "<|reserved_special_token_204|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128210": {
+      "content": "<|reserved_special_token_205|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128211": {
+      "content": "<|reserved_special_token_206|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128212": {
+      "content": "<|reserved_special_token_207|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128213": {
+      "content": "<|reserved_special_token_208|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128214": {
+      "content": "<|reserved_special_token_209|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128215": {
+      "content": "<|reserved_special_token_210|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128216": {
+      "content": "<|reserved_special_token_211|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128217": {
+      "content": "<|reserved_special_token_212|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128218": {
+      "content": "<|reserved_special_token_213|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128219": {
+      "content": "<|reserved_special_token_214|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128220": {
+      "content": "<|reserved_special_token_215|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128221": {
+      "content": "<|reserved_special_token_216|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128222": {
+      "content": "<|reserved_special_token_217|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128223": {
+      "content": "<|reserved_special_token_218|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128224": {
+      "content": "<|reserved_special_token_219|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128225": {
+      "content": "<|reserved_special_token_220|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128226": {
+      "content": "<|reserved_special_token_221|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128227": {
+      "content": "<|reserved_special_token_222|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128228": {
+      "content": "<|reserved_special_token_223|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128229": {
+      "content": "<|reserved_special_token_224|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128230": {
+      "content": "<|reserved_special_token_225|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128231": {
+      "content": "<|reserved_special_token_226|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128232": {
+      "content": "<|reserved_special_token_227|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128233": {
+      "content": "<|reserved_special_token_228|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128234": {
+      "content": "<|reserved_special_token_229|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128235": {
+      "content": "<|reserved_special_token_230|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128236": {
+      "content": "<|reserved_special_token_231|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128237": {
+      "content": "<|reserved_special_token_232|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128238": {
+      "content": "<|reserved_special_token_233|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128239": {
+      "content": "<|reserved_special_token_234|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128240": {
+      "content": "<|reserved_special_token_235|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128241": {
+      "content": "<|reserved_special_token_236|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128242": {
+      "content": "<|reserved_special_token_237|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128243": {
+      "content": "<|reserved_special_token_238|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128244": {
+      "content": "<|reserved_special_token_239|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128245": {
+      "content": "<|reserved_special_token_240|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128246": {
+      "content": "<|reserved_special_token_241|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128247": {
+      "content": "<|reserved_special_token_242|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128248": {
+      "content": "<|reserved_special_token_243|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128249": {
+      "content": "<|reserved_special_token_244|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128250": {
+      "content": "<|reserved_special_token_245|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128251": {
+      "content": "<|reserved_special_token_246|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128252": {
+      "content": "<|reserved_special_token_247|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128253": {
+      "content": "<|reserved_special_token_248|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128254": {
+      "content": "<|reserved_special_token_249|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128255": {
+      "content": "<|reserved_special_token_250|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128256": {
+      "content": "<SEP>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "extra_special_tokens": {},
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|end_of_text|>",
+  "sep_token": "<SEP>",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

trainer_state.json

@@ -0,0 +1,1258 @@
+{
+  "best_global_step": 50000,
+  "best_metric": 0.5881,
+  "best_model_checkpoint": "models/NED/SPACCC_full_upsampled_tfidf/Meta-Llama-3-8B-Instruct/checkpoint-50000",
+  "epoch": 3.0,
+  "eval_steps": 2000,
+  "global_step": 103965,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "entropy": 1.771104082763195,
+      "epoch": 0.057711729909104025,
+      "grad_norm": 8.875,
+      "learning_rate": 1.9227316447579353e-05,
+      "loss": 0.807,
+      "mean_token_accuracy": 0.8383643639683723,
+      "num_tokens": 15534479.0,
+      "step": 2000
+    },
+    {
+      "epoch": 0.057711729909104025,
+      "eval_entropy": 1.7064778925563795,
+      "eval_loss": 0.5752137899398804,
+      "eval_mean_token_accuracy": 0.8736604764762166,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 15534479.0,
+      "eval_recall": 0.4857,
+      "eval_runtime": 19.5541,
+      "eval_samples_per_second": 46.435,
+      "eval_steps_per_second": 46.435,
+      "step": 2000
+    },
+    {
+      "entropy": 1.7855026668310165,
+      "epoch": 0.11542345981820805,
+      "grad_norm": 4.9375,
+      "learning_rate": 2.973821470360748e-05,
+      "loss": 0.4284,
+      "mean_token_accuracy": 0.9000137696564198,
+      "num_tokens": 31091376.0,
+      "step": 4000
+    },
+    {
+      "epoch": 0.11542345981820805,
+      "eval_entropy": 1.9107416086391205,
+      "eval_loss": 0.6141767501831055,
+      "eval_mean_token_accuracy": 0.8813558152921924,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 31091376.0,
+      "eval_recall": 0.5132,
+      "eval_runtime": 18.6743,
+      "eval_samples_per_second": 48.623,
+      "eval_steps_per_second": 48.623,
+      "step": 4000
+    },
+    {
+      "entropy": 1.7454647228717803,
+      "epoch": 0.17313518972731207,
+      "grad_norm": 6.40625,
+      "learning_rate": 2.914324812089721e-05,
+      "loss": 0.3478,
+      "mean_token_accuracy": 0.9183464118242264,
+      "num_tokens": 46695519.0,
+      "step": 6000
+    },
+    {
+      "epoch": 0.17313518972731207,
+      "eval_entropy": 1.6889207995279245,
+      "eval_loss": 0.7172139286994934,
+      "eval_mean_token_accuracy": 0.8881947658627831,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 46695519.0,
+      "eval_recall": 0.5551,
+      "eval_runtime": 18.7505,
+      "eval_samples_per_second": 48.425,
+      "eval_steps_per_second": 48.425,
+      "step": 6000
+    },
+    {
+      "entropy": 1.686599359869957,
+      "epoch": 0.2308469196364161,
+      "grad_norm": 5.0,
+      "learning_rate": 2.8548281538186937e-05,
+      "loss": 0.2911,
+      "mean_token_accuracy": 0.9297748121023178,
+      "num_tokens": 62201311.0,
+      "step": 8000
+    },
+    {
+      "epoch": 0.2308469196364161,
+      "eval_entropy": 1.765830583837589,
+      "eval_loss": 0.6870580315589905,
+      "eval_mean_token_accuracy": 0.8954090231184392,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 62201311.0,
+      "eval_recall": 0.554,
+      "eval_runtime": 18.2183,
+      "eval_samples_per_second": 49.84,
+      "eval_steps_per_second": 49.84,
+      "step": 8000
+    },
+    {
+      "entropy": 1.6513705806136132,
+      "epoch": 0.28855864954552013,
+      "grad_norm": 2.890625,
+      "learning_rate": 2.795331495547667e-05,
+      "loss": 0.2723,
+      "mean_token_accuracy": 0.9330688781142235,
+      "num_tokens": 77723725.0,
+      "step": 10000
+    },
+    {
+      "epoch": 0.28855864954552013,
+      "eval_entropy": 1.7299244680199855,
+      "eval_loss": 0.671196460723877,
+      "eval_mean_token_accuracy": 0.8996821519003828,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 77723725.0,
+      "eval_recall": 0.5837,
+      "eval_runtime": 18.1693,
+      "eval_samples_per_second": 49.975,
+      "eval_steps_per_second": 49.975,
+      "step": 10000
+    },
+    {
+      "entropy": 1.6137347612977029,
+      "epoch": 0.34627037945462413,
+      "grad_norm": 4.15625,
+      "learning_rate": 2.7358348372766396e-05,
+      "loss": 0.247,
+      "mean_token_accuracy": 0.938195524007082,
+      "num_tokens": 93268948.0,
+      "step": 12000
+    },
+    {
+      "epoch": 0.34627037945462413,
+      "eval_entropy": 1.642808350631844,
+      "eval_loss": 0.7069945335388184,
+      "eval_mean_token_accuracy": 0.9005183333998735,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 93268948.0,
+      "eval_recall": 0.5683,
+      "eval_runtime": 18.2098,
+      "eval_samples_per_second": 49.863,
+      "eval_steps_per_second": 49.863,
+      "step": 12000
+    },
+    {
+      "entropy": 1.5417808018922805,
+      "epoch": 0.4039821093637282,
+      "grad_norm": 6.15625,
+      "learning_rate": 2.6763381790056127e-05,
+      "loss": 0.2369,
+      "mean_token_accuracy": 0.9400165711343288,
+      "num_tokens": 108712209.0,
+      "step": 14000
+    },
+    {
+      "epoch": 0.4039821093637282,
+      "eval_entropy": 1.6070995822483223,
+      "eval_loss": 0.7336843013763428,
+      "eval_mean_token_accuracy": 0.9035924484026064,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 108712209.0,
+      "eval_recall": 0.5749,
+      "eval_runtime": 18.1147,
+      "eval_samples_per_second": 50.125,
+      "eval_steps_per_second": 50.125,
+      "step": 14000
+    },
+    {
+      "entropy": 1.5307720832824707,
+      "epoch": 0.4616938392728322,
+      "grad_norm": 1.3203125,
+      "learning_rate": 2.6168415207345855e-05,
+      "loss": 0.224,
+      "mean_token_accuracy": 0.9427571404874325,
+      "num_tokens": 124237300.0,
+      "step": 16000
+    },
+    {
+      "epoch": 0.4616938392728322,
+      "eval_entropy": 1.6007142412242386,
+      "eval_loss": 0.723686933517456,
+      "eval_mean_token_accuracy": 0.9036177859891878,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 124237300.0,
+      "eval_recall": 0.5738,
+      "eval_runtime": 18.0095,
+      "eval_samples_per_second": 50.418,
+      "eval_steps_per_second": 50.418,
+      "step": 16000
+    },
+    {
+      "entropy": 1.5197161840200424,
+      "epoch": 0.5194055691819363,
+      "grad_norm": 4.375,
+      "learning_rate": 2.5573448624635583e-05,
+      "loss": 0.2127,
+      "mean_token_accuracy": 0.9449084457457065,
+      "num_tokens": 139804083.0,
+      "step": 18000
+    },
+    {
+      "epoch": 0.5194055691819363,
+      "eval_entropy": 1.5377285603909765,
+      "eval_loss": 0.7464824318885803,
+      "eval_mean_token_accuracy": 0.9009800682968505,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 139804083.0,
+      "eval_recall": 0.5694,
+      "eval_runtime": 17.3142,
+      "eval_samples_per_second": 52.443,
+      "eval_steps_per_second": 52.443,
+      "step": 18000
+    },
+    {
+      "entropy": 1.4599213127493857,
+      "epoch": 0.5771172990910403,
+      "grad_norm": 4.96875,
+      "learning_rate": 2.4978482041925314e-05,
+      "loss": 0.2026,
+      "mean_token_accuracy": 0.9471077627837657,
+      "num_tokens": 155379953.0,
+      "step": 20000
+    },
+    {
+      "epoch": 0.5771172990910403,
+      "eval_entropy": 1.475688523001608,
+      "eval_loss": 0.762593150138855,
+      "eval_mean_token_accuracy": 0.9033181490089399,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 155379953.0,
+      "eval_recall": 0.5705,
+      "eval_runtime": 17.8659,
+      "eval_samples_per_second": 50.823,
+      "eval_steps_per_second": 50.823,
+      "step": 20000
+    },
+    {
+      "entropy": 1.4277473657727242,
+      "epoch": 0.6348290290001443,
+      "grad_norm": 5.15625,
+      "learning_rate": 2.4383515459215042e-05,
+      "loss": 0.1907,
+      "mean_token_accuracy": 0.9498722539842128,
+      "num_tokens": 170982670.0,
+      "step": 22000
+    },
+    {
+      "epoch": 0.6348290290001443,
+      "eval_entropy": 1.5171283333837198,
+      "eval_loss": 0.7434535026550293,
+      "eval_mean_token_accuracy": 0.9021543205965983,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 170982670.0,
+      "eval_recall": 0.5727,
+      "eval_runtime": 17.4411,
+      "eval_samples_per_second": 52.061,
+      "eval_steps_per_second": 52.061,
+      "step": 22000
+    },
+    {
+      "entropy": 1.4187169399261474,
+      "epoch": 0.6925407589092483,
+      "grad_norm": 2.9375,
+      "learning_rate": 2.378854887650477e-05,
+      "loss": 0.1905,
+      "mean_token_accuracy": 0.9493927232325077,
+      "num_tokens": 186663193.0,
+      "step": 24000
+    },
+    {
+      "epoch": 0.6925407589092483,
+      "eval_entropy": 1.4800235901503835,
+      "eval_loss": 0.7626135945320129,
+      "eval_mean_token_accuracy": 0.9030716799298047,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 186663193.0,
+      "eval_recall": 0.5771,
+      "eval_runtime": 17.6742,
+      "eval_samples_per_second": 51.374,
+      "eval_steps_per_second": 51.374,
+      "step": 24000
+    },
+    {
+      "entropy": 1.3858990859389304,
+      "epoch": 0.7502524888183524,
+      "grad_norm": 5.4375,
+      "learning_rate": 2.31935822937945e-05,
+      "loss": 0.1847,
+      "mean_token_accuracy": 0.9506744608581066,
+      "num_tokens": 202382574.0,
+      "step": 26000
+    },
+    {
+      "epoch": 0.7502524888183524,
+      "eval_entropy": 1.4523340987476483,
+      "eval_loss": 0.7447758316993713,
+      "eval_mean_token_accuracy": 0.9053517018777159,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 202382574.0,
+      "eval_recall": 0.5837,
+      "eval_runtime": 17.5503,
+      "eval_samples_per_second": 51.737,
+      "eval_steps_per_second": 51.737,
+      "step": 26000
+    },
+    {
+      "entropy": 1.371607663989067,
+      "epoch": 0.8079642187274564,
+      "grad_norm": 6.46875,
+      "learning_rate": 2.259861571108423e-05,
+      "loss": 0.1747,
+      "mean_token_accuracy": 0.9534294557571411,
+      "num_tokens": 217894933.0,
+      "step": 28000
+    },
+    {
+      "epoch": 0.8079642187274564,
+      "eval_entropy": 1.449135780728336,
+      "eval_loss": 0.726492166519165,
+      "eval_mean_token_accuracy": 0.9063065528607054,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 217894933.0,
+      "eval_recall": 0.5782,
+      "eval_runtime": 17.4046,
+      "eval_samples_per_second": 52.17,
+      "eval_steps_per_second": 52.17,
+      "step": 28000
+    },
+    {
+      "entropy": 1.3474767149090767,
+      "epoch": 0.8656759486365604,
+      "grad_norm": 1.15625,
+      "learning_rate": 2.2003649128373957e-05,
+      "loss": 0.1662,
+      "mean_token_accuracy": 0.9552805411219597,
+      "num_tokens": 233442927.0,
+      "step": 30000
+    },
+    {
+      "epoch": 0.8656759486365604,
+      "eval_entropy": 1.4218564205495272,
+      "eval_loss": 0.7416213750839233,
+      "eval_mean_token_accuracy": 0.9069001922129535,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 233442927.0,
+      "eval_recall": 0.587,
+      "eval_runtime": 17.6827,
+      "eval_samples_per_second": 51.35,
+      "eval_steps_per_second": 51.35,
+      "step": 30000
+    },
+    {
+      "entropy": 1.337467650592327,
+      "epoch": 0.9233876785456644,
+      "grad_norm": 4.375,
+      "learning_rate": 2.1408682545663684e-05,
+      "loss": 0.1637,
+      "mean_token_accuracy": 0.9557446602284908,
+      "num_tokens": 249071659.0,
+      "step": 32000
+    },
+    {
+      "epoch": 0.9233876785456644,
+      "eval_entropy": 1.3889587524835234,
+      "eval_loss": 0.7390624284744263,
+      "eval_mean_token_accuracy": 0.9073679444876537,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 249071659.0,
+      "eval_recall": 0.5848,
+      "eval_runtime": 17.5032,
+      "eval_samples_per_second": 51.876,
+      "eval_steps_per_second": 51.876,
+      "step": 32000
+    },
+    {
+      "entropy": 1.3001156712770463,
+      "epoch": 0.9810994084547684,
+      "grad_norm": 3.890625,
+      "learning_rate": 2.0813715962953412e-05,
+      "loss": 0.1557,
+      "mean_token_accuracy": 0.9573536138236522,
+      "num_tokens": 264672169.0,
+      "step": 34000
+    },
+    {
+      "epoch": 0.9810994084547684,
+      "eval_entropy": 1.4234498099214705,
+      "eval_loss": 0.7484801411628723,
+      "eval_mean_token_accuracy": 0.9075630426012997,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 264672169.0,
+      "eval_recall": 0.5859,
+      "eval_runtime": 16.2686,
+      "eval_samples_per_second": 55.813,
+      "eval_steps_per_second": 55.813,
+      "step": 34000
+    },
+    {
+      "entropy": 1.2362971892952919,
+      "epoch": 1.0388111383638725,
+      "grad_norm": 4.96875,
+      "learning_rate": 2.0218749380243143e-05,
+      "loss": 0.1179,
+      "mean_token_accuracy": 0.9659893708825111,
+      "num_tokens": 280288735.0,
+      "step": 36000
+    },
+    {
+      "epoch": 1.0388111383638725,
+      "eval_entropy": 1.268248316540592,
+      "eval_loss": 0.7719414830207825,
+      "eval_mean_token_accuracy": 0.9055645169796923,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 280288735.0,
+      "eval_recall": 0.5771,
+      "eval_runtime": 16.2819,
+      "eval_samples_per_second": 55.767,
+      "eval_steps_per_second": 55.767,
+      "step": 36000
+    },
+    {
+      "entropy": 1.1861163977086544,
+      "epoch": 1.0965228682729764,
+      "grad_norm": 2.703125,
+      "learning_rate": 1.962378279753287e-05,
+      "loss": 0.1001,
+      "mean_token_accuracy": 0.9700573923885822,
+      "num_tokens": 295889252.0,
+      "step": 38000
+    },
+    {
+      "epoch": 1.0965228682729764,
+      "eval_entropy": 1.2859099207339308,
+      "eval_loss": 0.7740228176116943,
+      "eval_mean_token_accuracy": 0.9053704237425905,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 295889252.0,
+      "eval_recall": 0.5815,
+      "eval_runtime": 17.2227,
+      "eval_samples_per_second": 52.721,
+      "eval_steps_per_second": 52.721,
+      "step": 38000
+    },
+    {
+      "entropy": 1.172551353752613,
+      "epoch": 1.1542345981820805,
+      "grad_norm": 5.09375,
+      "learning_rate": 1.90288162148226e-05,
+      "loss": 0.0999,
+      "mean_token_accuracy": 0.9701334120929241,
+      "num_tokens": 311460812.0,
+      "step": 40000
+    },
+    {
+      "epoch": 1.1542345981820805,
+      "eval_entropy": 1.2553664291613953,
+      "eval_loss": 0.7721803784370422,
+      "eval_mean_token_accuracy": 0.9035136323537071,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 311460812.0,
+      "eval_recall": 0.576,
+      "eval_runtime": 16.7852,
+      "eval_samples_per_second": 54.095,
+      "eval_steps_per_second": 54.095,
+      "step": 40000
+    },
+    {
+      "entropy": 1.1619984501898288,
+      "epoch": 1.2119463280911846,
+      "grad_norm": 2.21875,
+      "learning_rate": 1.843384963211233e-05,
+      "loss": 0.0965,
+      "mean_token_accuracy": 0.9710877353549003,
+      "num_tokens": 327135410.0,
+      "step": 42000
+    },
+    {
+      "epoch": 1.2119463280911846,
+      "eval_entropy": 1.2579897131426219,
+      "eval_loss": 0.7692885398864746,
+      "eval_mean_token_accuracy": 0.9031982754033042,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 327135410.0,
+      "eval_recall": 0.576,
+      "eval_runtime": 17.2218,
+      "eval_samples_per_second": 52.724,
+      "eval_steps_per_second": 52.724,
+      "step": 42000
+    },
+    {
+      "entropy": 1.1411213338077069,
+      "epoch": 1.2696580580002885,
+      "grad_norm": 3.578125,
+      "learning_rate": 1.7838883049402058e-05,
+      "loss": 0.0944,
+      "mean_token_accuracy": 0.9715765230953693,
+      "num_tokens": 342774177.0,
+      "step": 44000
+    },
+    {
+      "epoch": 1.2696580580002885,
+      "eval_entropy": 1.253242742319464,
+      "eval_loss": 0.7519774436950684,
+      "eval_mean_token_accuracy": 0.9045354708175827,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 342774177.0,
+      "eval_recall": 0.576,
+      "eval_runtime": 17.1793,
+      "eval_samples_per_second": 52.854,
+      "eval_steps_per_second": 52.854,
+      "step": 44000
+    },
+    {
+      "entropy": 1.1595853001475334,
+      "epoch": 1.3273697879093926,
+      "grad_norm": 5.5,
+      "learning_rate": 1.724391646669179e-05,
+      "loss": 0.0963,
+      "mean_token_accuracy": 0.9711391851603984,
+      "num_tokens": 358312922.0,
+      "step": 46000
+    },
+    {
+      "epoch": 1.3273697879093926,
+      "eval_entropy": 1.2453804171689282,
+      "eval_loss": 0.7676454186439514,
+      "eval_mean_token_accuracy": 0.9064169454876547,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 358312922.0,
+      "eval_recall": 0.5837,
+      "eval_runtime": 16.8037,
+      "eval_samples_per_second": 54.036,
+      "eval_steps_per_second": 54.036,
+      "step": 46000
+    },
+    {
+      "entropy": 1.1609133576154709,
+      "epoch": 1.3850815178184965,
+      "grad_norm": 4.03125,
+      "learning_rate": 1.6648949883981517e-05,
+      "loss": 0.0922,
+      "mean_token_accuracy": 0.9723608312606812,
+      "num_tokens": 373752333.0,
+      "step": 48000
+    },
+    {
+      "epoch": 1.3850815178184965,
+      "eval_entropy": 1.2345776492934921,
+      "eval_loss": 0.7665285468101501,
+      "eval_mean_token_accuracy": 0.9063460667687365,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 373752333.0,
+      "eval_recall": 0.5793,
+      "eval_runtime": 16.7578,
+      "eval_samples_per_second": 54.184,
+      "eval_steps_per_second": 54.184,
+      "step": 48000
+    },
+    {
+      "entropy": 1.1655547478497028,
+      "epoch": 1.4427932477276006,
+      "grad_norm": 5.59375,
+      "learning_rate": 1.6053983301271245e-05,
+      "loss": 0.094,
+      "mean_token_accuracy": 0.9717481000125409,
+      "num_tokens": 389447345.0,
+      "step": 50000
+    },
+    {
+      "epoch": 1.4427932477276006,
+      "eval_entropy": 1.2292915042407713,
+      "eval_loss": 0.7735024094581604,
+      "eval_mean_token_accuracy": 0.907910385517822,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 389447345.0,
+      "eval_recall": 0.5881,
+      "eval_runtime": 17.3697,
+      "eval_samples_per_second": 52.275,
+      "eval_steps_per_second": 52.275,
+      "step": 50000
+    },
+    {
+      "entropy": 1.1552352701127528,
+      "epoch": 1.5005049776367048,
+      "grad_norm": 3.0,
+      "learning_rate": 1.5459016718560976e-05,
+      "loss": 0.091,
+      "mean_token_accuracy": 0.9726284679472447,
+      "num_tokens": 404935652.0,
+      "step": 52000
+    },
+    {
+      "epoch": 1.5005049776367048,
+      "eval_entropy": 1.2490241264325406,
+      "eval_loss": 0.7779573202133179,
+      "eval_mean_token_accuracy": 0.9046718338053132,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 404935652.0,
+      "eval_recall": 0.5848,
+      "eval_runtime": 17.2173,
+      "eval_samples_per_second": 52.738,
+      "eval_steps_per_second": 52.738,
+      "step": 52000
+    },
+    {
+      "entropy": 1.1580015743076801,
+      "epoch": 1.5582167075458087,
+      "grad_norm": 0.0019989013671875,
+      "learning_rate": 1.4864050135850704e-05,
+      "loss": 0.0674,
+      "mean_token_accuracy": 0.979576114565134,
+      "num_tokens": 15533221.0,
+      "step": 54000
+    },
+    {
+      "epoch": 1.5582167075458087,
+      "eval_entropy": 1.204221866138706,
+      "eval_loss": 0.8085830211639404,
+      "eval_mean_token_accuracy": 0.9048162211668124,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 15533221.0,
+      "eval_recall": 0.5738,
+      "eval_runtime": 17.1435,
+      "eval_samples_per_second": 52.965,
+      "eval_steps_per_second": 52.965,
+      "step": 54000
+    },
+    {
+      "entropy": 1.144241349697113,
+      "epoch": 1.6159284374549125,
+      "grad_norm": 3.765625,
+      "learning_rate": 1.4269083553140432e-05,
+      "loss": 0.0633,
+      "mean_token_accuracy": 0.9807874869704246,
+      "num_tokens": 31150685.0,
+      "step": 56000
+    },
+    {
+      "epoch": 1.6159284374549125,
+      "eval_entropy": 1.2577752770306256,
+      "eval_loss": 0.8108322024345398,
+      "eval_mean_token_accuracy": 0.904205797002179,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 31150685.0,
+      "eval_recall": 0.5738,
+      "eval_runtime": 16.4274,
+      "eval_samples_per_second": 55.274,
+      "eval_steps_per_second": 55.274,
+      "step": 56000
+    },
+    {
+      "entropy": 1.162702257514,
+      "epoch": 1.6736401673640167,
+      "grad_norm": 2.8125,
+      "learning_rate": 1.3674116970430161e-05,
+      "loss": 0.0665,
+      "mean_token_accuracy": 0.9797295650243759,
+      "num_tokens": 46832332.0,
+      "step": 58000
+    },
+    {
+      "epoch": 1.6736401673640167,
+      "eval_entropy": 1.268515376989537,
+      "eval_loss": 0.814584493637085,
+      "eval_mean_token_accuracy": 0.904229478295154,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 46832332.0,
+      "eval_recall": 0.5716,
+      "eval_runtime": 16.4853,
+      "eval_samples_per_second": 55.079,
+      "eval_steps_per_second": 55.079,
+      "step": 58000
+    },
+    {
+      "entropy": 1.1661596206724645,
+      "epoch": 1.7313518972731208,
+      "grad_norm": 5.96875,
+      "learning_rate": 1.3079150387719889e-05,
+      "loss": 0.0672,
+      "mean_token_accuracy": 0.9796462517380714,
+      "num_tokens": 62558817.0,
+      "step": 60000
+    },
+    {
+      "epoch": 1.7313518972731208,
+      "eval_entropy": 1.2622852631996382,
+      "eval_loss": 0.8227198123931885,
+      "eval_mean_token_accuracy": 0.9038923141846048,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 62558817.0,
+      "eval_recall": 0.5716,
+      "eval_runtime": 16.9948,
+      "eval_samples_per_second": 53.428,
+      "eval_steps_per_second": 53.428,
+      "step": 60000
+    },
+    {
+      "entropy": 1.1797457176148891,
+      "epoch": 1.789063627182225,
+      "grad_norm": 0.5546875,
+      "learning_rate": 1.2484183805009618e-05,
+      "loss": 0.0657,
+      "mean_token_accuracy": 0.980204150468111,
+      "num_tokens": 78074806.0,
+      "step": 62000
+    },
+    {
+      "epoch": 1.789063627182225,
+      "eval_entropy": 1.2418163208052975,
+      "eval_loss": 0.8185028433799744,
+      "eval_mean_token_accuracy": 0.9041991046740621,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 78074806.0,
+      "eval_recall": 0.5705,
+      "eval_runtime": 17.1144,
+      "eval_samples_per_second": 53.055,
+      "eval_steps_per_second": 53.055,
+      "step": 62000
+    },
+    {
+      "entropy": 1.1774089051187038,
+      "epoch": 1.8467753570913288,
+      "grad_norm": 5.5625,
+      "learning_rate": 1.1889217222299348e-05,
+      "loss": 0.0649,
+      "mean_token_accuracy": 0.9804997465908527,
+      "num_tokens": 93602629.0,
+      "step": 64000
+    },
+    {
+      "epoch": 1.8467753570913288,
+      "eval_entropy": 1.2988805646807087,
+      "eval_loss": 0.8260899782180786,
+      "eval_mean_token_accuracy": 0.9030656689523601,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 93602629.0,
+      "eval_recall": 0.576,
+      "eval_runtime": 16.1643,
+      "eval_samples_per_second": 56.173,
+      "eval_steps_per_second": 56.173,
+      "step": 64000
+    },
+    {
+      "entropy": 1.174987347126007,
+      "epoch": 1.9044870870004327,
+      "grad_norm": 5.53125,
+      "learning_rate": 1.1294250639589077e-05,
+      "loss": 0.064,
+      "mean_token_accuracy": 0.9806980607807636,
+      "num_tokens": 109249414.0,
+      "step": 66000
+    },
+    {
+      "epoch": 1.9044870870004327,
+      "eval_entropy": 1.2433809736489199,
+      "eval_loss": 0.8272661566734314,
+      "eval_mean_token_accuracy": 0.9028221254569319,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 109249414.0,
+      "eval_recall": 0.5727,
+      "eval_runtime": 16.3988,
+      "eval_samples_per_second": 55.37,
+      "eval_steps_per_second": 55.37,
+      "step": 66000
+    },
+    {
+      "entropy": 1.1633582679629326,
+      "epoch": 1.9621988169095368,
+      "grad_norm": 5.34375,
+      "learning_rate": 1.0699284056878807e-05,
+      "loss": 0.0643,
+      "mean_token_accuracy": 0.9805754337012768,
+      "num_tokens": 124880720.0,
+      "step": 68000
+    },
+    {
+      "epoch": 1.9621988169095368,
+      "eval_entropy": 1.224490842367584,
+      "eval_loss": 0.8288715481758118,
+      "eval_mean_token_accuracy": 0.9034351931991557,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 124880720.0,
+      "eval_recall": 0.5738,
+      "eval_runtime": 16.4997,
+      "eval_samples_per_second": 55.031,
+      "eval_steps_per_second": 55.031,
+      "step": 68000
+    },
+    {
+      "entropy": 1.1513627296090125,
+      "epoch": 2.019910546818641,
+      "grad_norm": 0.79296875,
+      "learning_rate": 1.0104317474168535e-05,
+      "loss": 0.0633,
+      "mean_token_accuracy": 0.9811660476624966,
+      "num_tokens": 140499220.0,
+      "step": 70000
+    },
+    {
+      "epoch": 2.019910546818641,
+      "eval_entropy": 1.2267822175561593,
+      "eval_loss": 0.8458257913589478,
+      "eval_mean_token_accuracy": 0.9038964834572986,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 140499220.0,
+      "eval_recall": 0.5705,
+      "eval_runtime": 16.4967,
+      "eval_samples_per_second": 55.041,
+      "eval_steps_per_second": 55.041,
+      "step": 70000
+    },
+    {
+      "entropy": 1.143776093840599,
+      "epoch": 2.077622276727745,
+      "grad_norm": 7.9375,
+      "learning_rate": 9.509350891458264e-06,
+      "loss": 0.0597,
+      "mean_token_accuracy": 0.9825106913745403,
+      "num_tokens": 156048918.0,
+      "step": 72000
+    },
+    {
+      "epoch": 2.077622276727745,
+      "eval_entropy": 1.2148328015195116,
+      "eval_loss": 0.8337165713310242,
+      "eval_mean_token_accuracy": 0.9035390550475814,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 156048918.0,
+      "eval_recall": 0.5705,
+      "eval_runtime": 16.506,
+      "eval_samples_per_second": 55.01,
+      "eval_steps_per_second": 55.01,
+      "step": 72000
+    },
+    {
+      "entropy": 1.1460822140574456,
+      "epoch": 2.135334006636849,
+      "grad_norm": 12.375,
+      "learning_rate": 8.914384308747992e-06,
+      "loss": 0.0596,
+      "mean_token_accuracy": 0.98244061678648,
+      "num_tokens": 171653895.0,
+      "step": 74000
+    },
+    {
+      "epoch": 2.135334006636849,
+      "eval_entropy": 1.2635613490175046,
+      "eval_loss": 0.8348618745803833,
+      "eval_mean_token_accuracy": 0.9038379774285308,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 171653895.0,
+      "eval_recall": 0.5694,
+      "eval_runtime": 16.4822,
+      "eval_samples_per_second": 55.09,
+      "eval_steps_per_second": 55.09,
+      "step": 74000
+    },
+    {
+      "entropy": 1.1560133908391,
+      "epoch": 2.193045736545953,
+      "grad_norm": 7.625,
+      "learning_rate": 8.319417726037721e-06,
+      "loss": 0.06,
+      "mean_token_accuracy": 0.9822552761137485,
+      "num_tokens": 187228261.0,
+      "step": 76000
+    },
+    {
+      "epoch": 2.193045736545953,
+      "eval_entropy": 1.2220293277554575,
+      "eval_loss": 0.8315507769584656,
+      "eval_mean_token_accuracy": 0.9036543207809263,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 187228261.0,
+      "eval_recall": 0.5716,
+      "eval_runtime": 16.703,
+      "eval_samples_per_second": 54.361,
+      "eval_steps_per_second": 54.361,
+      "step": 76000
+    },
+    {
+      "entropy": 1.1676500248610973,
+      "epoch": 2.250757466455057,
+      "grad_norm": 4.84375,
+      "learning_rate": 7.72445114332745e-06,
+      "loss": 0.0611,
+      "mean_token_accuracy": 0.9819406977891922,
+      "num_tokens": 202699683.0,
+      "step": 78000
+    },
+    {
+      "epoch": 2.250757466455057,
+      "eval_entropy": 1.2448954319638827,
+      "eval_loss": 0.8309385776519775,
+      "eval_mean_token_accuracy": 0.9030922418255113,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 202699683.0,
+      "eval_recall": 0.5694,
+      "eval_runtime": 16.603,
+      "eval_samples_per_second": 54.689,
+      "eval_steps_per_second": 54.689,
+      "step": 78000
+    },
+    {
+      "entropy": 1.1656713368594647,
+      "epoch": 2.308469196364161,
+      "grad_norm": 6.53125,
+      "learning_rate": 7.129484560617179e-06,
+      "loss": 0.0618,
+      "mean_token_accuracy": 0.9817487963140011,
+      "num_tokens": 218284466.0,
+      "step": 80000
+    },
+    {
+      "epoch": 2.308469196364161,
+      "eval_entropy": 1.255102663181952,
+      "eval_loss": 0.8435425162315369,
+      "eval_mean_token_accuracy": 0.902260869642974,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 218284466.0,
+      "eval_recall": 0.5661,
+      "eval_runtime": 16.8256,
+      "eval_samples_per_second": 53.965,
+      "eval_steps_per_second": 53.965,
+      "step": 80000
+    },
+    {
+      "entropy": 1.1597592905461789,
+      "epoch": 2.366180926273265,
+      "grad_norm": 2.5,
+      "learning_rate": 6.534517977906908e-06,
+      "loss": 0.0602,
+      "mean_token_accuracy": 0.9821576415896416,
+      "num_tokens": 233928452.0,
+      "step": 82000
+    },
+    {
+      "epoch": 2.366180926273265,
+      "eval_entropy": 1.2422783964924875,
+      "eval_loss": 0.8390738368034363,
+      "eval_mean_token_accuracy": 0.9032785006950605,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 233928452.0,
+      "eval_recall": 0.5683,
+      "eval_runtime": 16.7577,
+      "eval_samples_per_second": 54.184,
+      "eval_steps_per_second": 54.184,
+      "step": 82000
+    },
+    {
+      "entropy": 1.17008468157053,
+      "epoch": 2.4238926561823693,
+      "grad_norm": 0.0400390625,
+      "learning_rate": 5.939551395196637e-06,
+      "loss": 0.0591,
+      "mean_token_accuracy": 0.9825585896968841,
+      "num_tokens": 249419664.0,
+      "step": 84000
+    },
+    {
+      "epoch": 2.4238926561823693,
+      "eval_entropy": 1.2469606770424067,
+      "eval_loss": 0.8383654356002808,
+      "eval_mean_token_accuracy": 0.9040639832418921,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 249419664.0,
+      "eval_recall": 0.5705,
+      "eval_runtime": 16.6392,
+      "eval_samples_per_second": 54.57,
+      "eval_steps_per_second": 54.57,
+      "step": 84000
+    },
+    {
+      "entropy": 1.163529093414545,
+      "epoch": 2.481604386091473,
+      "grad_norm": 6.28125,
+      "learning_rate": 5.3445848124863655e-06,
+      "loss": 0.0568,
+      "mean_token_accuracy": 0.9832313210368157,
+      "num_tokens": 264982654.0,
+      "step": 86000
+    },
+    {
+      "epoch": 2.481604386091473,
+      "eval_entropy": 1.236849331908289,
+      "eval_loss": 0.8381890058517456,
+      "eval_mean_token_accuracy": 0.9027883698630438,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 264982654.0,
+      "eval_recall": 0.5672,
+      "eval_runtime": 16.5964,
+      "eval_samples_per_second": 54.711,
+      "eval_steps_per_second": 54.711,
+      "step": 86000
+    },
+    {
+      "entropy": 1.1701532056927682,
+      "epoch": 2.539316116000577,
+      "grad_norm": 0.87109375,
+      "learning_rate": 4.749618229776094e-06,
+      "loss": 0.0574,
+      "mean_token_accuracy": 0.9830155865848065,
+      "num_tokens": 280520807.0,
+      "step": 88000
+    },
+    {
+      "epoch": 2.539316116000577,
+      "eval_entropy": 1.2492524392673097,
+      "eval_loss": 0.839518666267395,
+      "eval_mean_token_accuracy": 0.9025986767681685,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 280520807.0,
+      "eval_recall": 0.5661,
+      "eval_runtime": 16.5519,
+      "eval_samples_per_second": 54.858,
+      "eval_steps_per_second": 54.858,
+      "step": 88000
+    },
+    {
+      "entropy": 1.167941878914833,
+      "epoch": 2.597027845909681,
+      "grad_norm": 0.451171875,
+      "learning_rate": 4.154651647065824e-06,
+      "loss": 0.0602,
+      "mean_token_accuracy": 0.9820802296400071,
+      "num_tokens": 296146535.0,
+      "step": 90000
+    },
+    {
+      "epoch": 2.597027845909681,
+      "eval_entropy": 1.2443812186234848,
+      "eval_loss": 0.8395401835441589,
+      "eval_mean_token_accuracy": 0.9034286766981764,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 296146535.0,
+      "eval_recall": 0.5683,
+      "eval_runtime": 16.788,
+      "eval_samples_per_second": 54.086,
+      "eval_steps_per_second": 54.086,
+      "step": 90000
+    },
+    {
+      "entropy": 1.1601335457861424,
+      "epoch": 2.6547395758187853,
+      "grad_norm": 4.6875,
+      "learning_rate": 3.559685064355552e-06,
+      "loss": 0.0584,
+      "mean_token_accuracy": 0.9827592859268188,
+      "num_tokens": 311778551.0,
+      "step": 92000
+    },
+    {
+      "epoch": 2.6547395758187853,
+      "eval_entropy": 1.2437387075324415,
+      "eval_loss": 0.836577296257019,
+      "eval_mean_token_accuracy": 0.9039776291419231,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 311778551.0,
+      "eval_recall": 0.5705,
+      "eval_runtime": 16.2272,
+      "eval_samples_per_second": 55.956,
+      "eval_steps_per_second": 55.956,
+      "step": 92000
+    },
+    {
+      "entropy": 1.1733056641221047,
+      "epoch": 2.712451305727889,
+      "grad_norm": 5.03125,
+      "learning_rate": 2.964718481645281e-06,
+      "loss": 0.0564,
+      "mean_token_accuracy": 0.9832823853492737,
+      "num_tokens": 327170479.0,
+      "step": 94000
+    },
+    {
+      "epoch": 2.712451305727889,
+      "eval_entropy": 1.2440849004337966,
+      "eval_loss": 0.8399211168289185,
+      "eval_mean_token_accuracy": 0.9033104040155326,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 327170479.0,
+      "eval_recall": 0.5683,
+      "eval_runtime": 16.2233,
+      "eval_samples_per_second": 55.969,
+      "eval_steps_per_second": 55.969,
+      "step": 94000
+    },
+    {
+      "entropy": 1.1586334483027458,
+      "epoch": 2.770163035636993,
+      "grad_norm": 3.953125,
+      "learning_rate": 2.36975189893501e-06,
+      "loss": 0.0585,
+      "mean_token_accuracy": 0.9826480825543403,
+      "num_tokens": 342791353.0,
+      "step": 96000
+    },
+    {
+      "epoch": 2.770163035636993,
+      "eval_entropy": 1.2412338042180444,
+      "eval_loss": 0.8378188610076904,
+      "eval_mean_token_accuracy": 0.9035194405572005,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 342791353.0,
+      "eval_recall": 0.5672,
+      "eval_runtime": 16.2077,
+      "eval_samples_per_second": 56.023,
+      "eval_steps_per_second": 56.023,
+      "step": 96000
+    },
+    {
+      "entropy": 1.1629991734027862,
+      "epoch": 2.827874765546097,
+      "grad_norm": 6.4375,
+      "learning_rate": 1.7747853162247388e-06,
+      "loss": 0.0608,
+      "mean_token_accuracy": 0.9821404512822628,
+      "num_tokens": 358436354.0,
+      "step": 98000
+    },
+    {
+      "epoch": 2.827874765546097,
+      "eval_entropy": 1.2435034370369848,
+      "eval_loss": 0.8380420207977295,
+      "eval_mean_token_accuracy": 0.9037704004327631,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 358436354.0,
+      "eval_recall": 0.5694,
+      "eval_runtime": 16.2198,
+      "eval_samples_per_second": 55.981,
+      "eval_steps_per_second": 55.981,
+      "step": 98000
+    },
+    {
+      "entropy": 1.1640874392092229,
+      "epoch": 2.8855864954552013,
+      "grad_norm": 0.8203125,
+      "learning_rate": 1.1798187335144677e-06,
+      "loss": 0.0574,
+      "mean_token_accuracy": 0.9829988768994808,
+      "num_tokens": 374029027.0,
+      "step": 100000
+    },
+    {
+      "epoch": 2.8855864954552013,
+      "eval_entropy": 1.2438825091207606,
+      "eval_loss": 0.8370459079742432,
+      "eval_mean_token_accuracy": 0.9030486140810445,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 374029027.0,
+      "eval_recall": 0.5661,
+      "eval_runtime": 17.3712,
+      "eval_samples_per_second": 52.271,
+      "eval_steps_per_second": 52.271,
+      "step": 100000
+    },
+    {
+      "entropy": 1.1683570961356162,
+      "epoch": 2.9432982253643054,
+      "grad_norm": 6.65625,
+      "learning_rate": 5.848521508041964e-07,
+      "loss": 0.0583,
+      "mean_token_accuracy": 0.982835016399622,
+      "num_tokens": 389554889.0,
+      "step": 102000
+    },
+    {
+      "epoch": 2.9432982253643054,
+      "eval_entropy": 1.2444104566710636,
+      "eval_loss": 0.8376456499099731,
+      "eval_mean_token_accuracy": 0.9037148623608282,
+      "eval_num_gold": 908,
+      "eval_num_guess": 908,
+      "eval_num_tokens": 389554889.0,
+      "eval_recall": 0.5683,
+      "eval_runtime": 16.2408,
+      "eval_samples_per_second": 55.909,
+      "eval_steps_per_second": 55.909,
+      "step": 102000
+    }
+  ],
+  "logging_steps": 2000,
+  "max_steps": 103965,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 3,
+  "save_steps": 2000,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 3.6466469785747587e+19,
+  "train_batch_size": 1,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

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