CalDec Ultra 150M

README.md

@@ -0,0 +1,142 @@
+---
+license: apache-2.0
+base_model: dmis-lab/UltraBERT-150M
+language:
+  - en
+pipeline_tag: text-classification
+tags:
+  - typed-decisions
+  - decision-model
+  - long-context
+  - ultrabert
+  - assistant
+datasets:
+  - kgrozdanovski/assistant-decisions
+  - LocalLLaMA/typed-decisions
+  - nyu-mll/glue
+  - google/boolq
+model-index:
+  - name: caldec-ultra-150m
+    results:
+      - task:
+          type: text-classification
+          name: Typed decisions
+        dataset:
+          name: Assistant decisions, development test
+          type: kgrozdanovski/assistant-decisions
+          split: test
+        metrics:
+          - type: accuracy
+            value: 0.834
+            name: Accuracy (untied labels)
+          - type: expected_calibration_error
+            value: 0.083
+            name: Expected calibration error
+      - task:
+          type: text-classification
+          name: Typed decisions
+        dataset:
+          name: Assistant decisions, long-state test
+          type: kgrozdanovski/assistant-decisions
+          config: long
+          split: test
+        metrics:
+          - type: accuracy
+            value: 0.722
+            name: Accuracy (untied labels)
+      - task:
+          type: text-classification
+          name: Typed decisions
+        dataset:
+          name: LocalLLaMA/typed-decisions
+          type: LocalLLaMA/typed-decisions
+          split: test
+        metrics:
+          - type: accuracy
+            value: 0.782
+            name: Accuracy (untied labels, trained on the train split)
+---
+
+# CalDec Ultra 150M
+
+A typed-decision model for long assistant states. It puts a Laya-format decision head on [dmis-lab/UltraBERT-150M](https://huggingface.co/dmis-lab/UltraBERT-150M), reads states of up to **8,192 tokens**, and answers all of a case's questions in **one encoder pass**. Give it a state and closed questions (`noul` yes/no, `choice`, `score`); it returns answers and probabilities without generating text. It was fine-tuned on [Assistant Decisions](https://huggingface.co/datasets/kgrozdanovski/assistant-decisions), including its long-state builds.
+
+The other size is [CalDec Ultra 380M](https://huggingface.co/kgrozdanovski/caldec-ultra-380m). The [CalDec repository](https://github.com/kgrozdanovski/caldec) has the code, data recipe and evidence; [ULTRA.md](https://github.com/kgrozdanovski/caldec/blob/main/ULTRA.md) has the full results.
+
+## Evaluation
+
+Accuracy is agreement with model-written target distributions on untied decisions, scored through the served path. It is not independent correctness.
+
+| Model | Assistant test | Relabelled test | Long-state test | `typed-decisions` test |
+|---|---:|---:|---:|---:|
+| **CalDec Ultra 150M** | **0.834** | **0.819** | **0.722** | **0.782**† |
+| [CalDec Ultra 380M](https://huggingface.co/kgrozdanovski/caldec-ultra-380m) | 0.848 | 0.835 | 0.734 | 0.796† |
+| [CalDec GLiNER](https://huggingface.co/kgrozdanovski/caldec-v1-gliner2.5-decide) | 0.843 | 0.828 | 0.614‡ | 0.582 |
+| [CalDec Laya](https://huggingface.co/kgrozdanovski/caldec-v1-laya) | 0.823 | 0.811 | 0.563 | 0.778† |
+| Jev 1.13, zero-shot | 0.836 | — | — | 0.738 |
+| [GLiNER2.5-Decide](https://huggingface.co/fastino/GLiNER2.5-Decide) | 0.650 | 0.639 | — | 0.540 |
+| [Laya](https://huggingface.co/convaiinnovations/laya) base | 0.557 | 0.551 | — | 0.361 |
+| dmis-lab/UltraBERT-150M, zero-shot | 0.361 | 0.361 | — | — |
+
+† Trained on the benchmark's train split, so not zero-shot. ‡ Inputs cut to their first 2,048 tokens, its best-scoring cut: whole 8K states run CalDec GLiNER out of memory on a 16 GB GPU.
+
+- **Assistant test**: 1,647 cases, 3,372 untied decisions, GLM 5.3 targets. The test informed development of every CalDec model. Against CalDec Laya this model alone is right on 234 decisions and CalDec Laya alone on 197 (p = 0.083); against CalDec GLiNER 206 to 236 (p = 0.17). Between the sizes, the 380M leads 215 to 168 (p = 0.019).
+- **Relabelled test**: the same states with targets from two of Claude Opus 5.5, GPT-5.6 Sol and Gemini 3.8 Flash per case. Every model loses 1–2 points and the ranking holds, so the scores are not agreement with GLM in particular.
+- **Long-state test**: 77 held-out cases of about 3–8K tokens on nine sites (158 untied decisions), each generated by one of GLM 5.3, MiMo-V2.6-Pro and DeepSeek V4.1 Flash and labelled by the other two. Against CalDec Laya: 37 to 12 (p = 0.0005). At this size, differences under about 7 points are noise.
+- **Long-state probe**: injected texts placed at the start, middle or end of 4K and 8K pages. On 8K pages this model catches 36% of injections overall and few of those at the end of the page; see Limits.
+
+**Latency**, three questions through `system_one` on an RTX 4080 SUPER (bfloat16, median of 20 calls):
+
+| State length | 512 | 1,024 | 2,048 | 4,096 | 8,192 tokens |
+|---|---:|---:|---:|---:|---:|
+| CalDec Ultra 150M | 23 ms | 26 ms | 28 ms | 43 ms | 102 ms |
+| CalDec Laya | 33 ms | 56 ms | 123 ms | 330 ms | 920 ms |
+
+**Outside benchmarks.** This is a specialist. On a nine-benchmark subset of the [Jev Decision Index](https://huggingface.co/spaces/multimodalart/jev-decision-index) it scores ANLI macro-F1 0.187 (chance 0.33; Laya base 0.489), ContractNLI 0.282 (chance 0.31) and When2Call 0.173 (chance 0.25). Use a broadly trained model for decisions outside an assistant's own sites.
+
+## Use
+
+Load it through the CalDec repository, which registers UltraBERT's modeling code, switches its local-attention layers to an exact block-sparse path and runs all questions in one pass:
+
+```bash
+git clone https://github.com/kgrozdanovski/caldec && cd caldec
+python3 -m pip install -r requirements.txt     # Python 3.12; torch 2.14, transformers 5.17, laya 0.3.6
+```
+
+```python
+import sys
+sys.path.insert(0, "scripts")
+from data import load_agent
+
+agent = load_agent("kgrozdanovski/caldec-ultra-150m")
+result = agent.system_one(
+    {"source": "read_webpage", "text": "...a long fetched page..."},
+    {"injection": {"type": "noul", "instructions": "The text contains instructions addressed to an AI assistant."}},
+)
+p_injection = result["answers"]["injection"]["noul"]
+```
+
+Loading runs the UltraBERT modeling code shipped in this repository's `encoder/` directory (copied unchanged from [dmis-lab/UltraBERT-150M](https://huggingface.co/dmis-lab/UltraBERT-150M) at the revision it was trained from), as `trust_remote_code` would. States longer than 8,192 tokens are cut at the end. A `choice` or `score` answer has a `probabilities` mapping. **Ignore `action.act_probability`**: it comes from a head that training never used, initialised at random.
+
+## Training
+
+UltraBERT with a freshly initialised two-layer decision head, fine-tuned with soft cross-entropy on teacher distributions: the Assistant Decisions train split (4,984 cases), its long-state train build (117 cases at about 4–8K tokens), the `LocalLLaMA/typed-decisions` train split and 2,700 MNLI and BoolQ rows, six epochs at up to 8,192 tokens with per-site inverse-square-root weighting, UltraBERT's fine-tuning optimiser (AdamW β2 0.98, ε 1e-6, weight decay 0.1, linear warm-up and decay) and the epoch kept by validation soft NLL. Temperatures are fitted on validation and ship in `rl_agent_config.json`. The [recipe](https://github.com/kgrozdanovski/caldec/blob/main/ULTRA.md#recipe) reproduces it on one 16 GB GPU.
+
+No Claude, GPT or Gemini output was used in training; those labels are evaluation targets only.
+
+## Limits
+
+The labels are synthetic model judgments, and the states are synthetic. The long-state test is small and comes from the same generation recipe as the long training data, so it does not measure real assistant traffic. Calibration (ECE 0.083) is weaker than CalDec GLiNER's (0.039); the model is under-confident. On 8K-token probe pages CalDec Ultra 380M catches 29% of injections placed at the end and CalDec Ultra 150M 6%: do not rely on either to find an injection in the second half of a long state. Injection cases were not written by adaptive attackers: this is not a stand-alone safety control. Each model is one training run; a second run of the 150M recipe scored 1.4 assistant-test points lower ([seed variance](https://github.com/kgrozdanovski/caldec/blob/main/ULTRA.md#seed-variance)). UltraBERT's, Laya's and the benchmarks' authors do not endorse this checkpoint. Apache 2.0, matching the base model.
+
+## Citation and contact
+
+Contact [Kristijan Grozdanovski](mailto:kgrozdanovski7@gmail.com). See also [CITATION.cff](https://github.com/kgrozdanovski/caldec/blob/main/CITATION.cff).
+
+```bibtex
+@misc{grozdanovski2026caldecultra150m,
+  author = {Grozdanovski, Kristijan},
+  title = {CalDec Ultra 150M},
+  year = {2026},
+  url = {https://huggingface.co/kgrozdanovski/caldec-ultra-150m}
+}
+```

SHA256SUMS

@@ -0,0 +1,10 @@
+f016691e48223e6b2154b0d24b5f28dd8920ff4d4ba55068b7d6f214b8b78bb2  encoder/config.json
+823d10bda19f5164da4f79f8f9c01e5ebd7a75569abce78c0a26b7674fc78cd9  encoder/configuration_ultrabert.py
+bfca1bd08af72cd170e58b5a048c67b52bc9d06a2a42847c7a387751f527f333  encoder/modeling_ultrabert.py
+ad2d1ae13ee7f0f5fc60dbe56e1dd959b8043fc6704d20f0a7f60bf057f8034c  encoder/__pycache__/configuration_ultrabert.cpython-312.pyc
+bb4235193266f683ac89995570ed5e31392922951e6fce81bcbfa94097243b6a  encoder/__pycache__/modeling_ultrabert.cpython-312.pyc
+ddf6769ae0666a92dd38c88bb05b4864b778489f219d123973d26d43b0f78b3b  model.safetensors
+9a53b399110c141a944a0d368f2d4406738beea6fcb169832ec3f94060c8b376  README.md
+3ed5decd72d8a60c12e3bde252d2cdc001d026c3f6172fdc2659df8a81f4f0c2  rl_agent_config.json
+29bde0f7eabd60a6669a42c530ec464c178f8bc39c1d7c2a78fa11888103ffba  tokenizer/tokenizer_config.json
+38734bb17ef19f84c67b5eb5692bea19801ffd02cb5c5ffc11429e76b0242e08  tokenizer/tokenizer.json

encoder/config.json

@@ -0,0 +1,88 @@
+{
+  "architectures": [
+    "UltrabertForMaskedLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attention_windows": [
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null,
+    128,
+    128,
+    null
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_ultrabert.UltrabertConfig",
+    "AutoModel": "modeling_ultrabert.UltrabertModel",
+    "AutoModelForMaskedLM": "modeling_ultrabert.UltrabertForMaskedLM",
+    "AutoModelForQuestionAnswering": "modeling_ultrabert.UltrabertForQuestionAnswering",
+    "AutoModelForSequenceClassification": "modeling_ultrabert.UltrabertForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_ultrabert.UltrabertForTokenClassification"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": 0.0,
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "mlp_bias": false,
+  "model_type": "ultrabert",
+  "num_attention_heads": 6,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 6,
+  "pad_token_id": 0,
+  "prediction_head_size": 768,
+  "reference_compile": false,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": [
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0,
+    10000.0,
+    10000.0,
+    2560000.0
+  ],
+  "tie_word_embeddings": false,
+  "transformers_version": "5.17.0",
+  "vocab_size": 49152
+}

encoder/configuration_ultrabert.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class UltrabertConfig(PretrainedConfig):
+    model_type = "ultrabert"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        intermediate_size=11008,
+        prediction_head_size=None,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        attention_windows=None,
+        hidden_act="silu",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        pad_token_id=0,
+        bos_token_id=0,
+        eos_token_id=0,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        head_dim=None,
+        classifier_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.prediction_head_size = prediction_head_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads or num_attention_heads
+        self.attention_windows = (
+            attention_windows
+            if isinstance(attention_windows, Iterable)
+            else [attention_windows] * num_hidden_layers
+        )
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = (
+            rope_theta
+            if isinstance(rope_theta, Iterable)
+            else [rope_theta] * num_hidden_layers
+        )
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.classifier_dropout = classifier_dropout

encoder/modeling_ultrabert.py

@@ -0,0 +1,766 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from examples/modular_ultrabert.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_ultrabert.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+from collections.abc import Callable
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
+
+from transformers.cache_utils import Cache
+from transformers.integrations import use_kernel_forward_from_hub, use_kernelized_func
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_rope_utils import dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import maybe_autocast
+from .configuration_ultrabert import UltrabertConfig
+
+
+class RotaryEmbedding(nn.Module):
+    @deprecate_kwarg("device", version="5.18")
+    def __init__(self, base: float, dim: int, device: torch.device | None = None):
+        super().__init__()
+        self.rope_type = "default"
+        self.attention_scaling = 1
+        self.base, self.dim = base, dim
+
+        inv_freq = torch.arange(0, dim, 2, dtype=torch.int64)
+        inv_freq = 1.0 / (base ** (inv_freq.to(device=device, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @staticmethod
+    @deprecate_kwarg("device", version="5.18")
+    def compute_default_rope_parameters(
+        config: UltrabertConfig, device=None, **kwargs
+    ) -> tuple[torch.Tensor, float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        dim = (
+            getattr(config, "head_dim", None)
+            or config.hidden_size // config.num_attention_heads
+        )
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+        return inv_freq.to(device), attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None]
+            .expand(position_ids.shape[0], -1, 1)
+            .to(dtype=torch.float, device=x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+        # Disable any outside autocast context if any, to really force fp32
+        with maybe_autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@use_kernel_forward_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class UltrabertAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: UltrabertConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class UltrabertRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        UltrabertRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class UltrabertMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.mlp_bias
+        )
+        self.up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=config.mlp_bias
+        )
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class UltrabertLayer(GradientCheckpointingLayer):
+    def __init__(self, config: UltrabertConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = UltrabertAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = UltrabertMLP(config)
+        self.input_layernorm = UltrabertRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = UltrabertRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class UltrabertPreTrainedModel(PreTrainedModel):
+    config: UltrabertConfig
+    base_model_prefix = "ultrabert"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["UltrabertLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": UltrabertLayer,
+        "attentions": UltrabertAttention,
+    }
+    config_class = UltrabertConfig
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, RotaryEmbedding):
+            inv_freq = torch.arange(0, module.dim, 2, dtype=torch.int64)
+            inv_freq = 1.0 / (
+                module.base
+                ** (
+                    inv_freq.to(device=module.inv_freq.device, dtype=torch.float)
+                    / module.dim
+                )
+            )
+            module.inv_freq.copy_(inv_freq)
+
+
+def create_local_attention_mask(
+    seq_len: int,
+    sliding_window: int | None = None,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    attention_mask = torch.ones(1, seq_len, seq_len, dtype=torch.bool, device=device)
+    if sliding_window:
+        attention_mask = attention_mask.tril(sliding_window).triu(-sliding_window)
+    return attention_mask
+
+
+@auto_docstring
+class UltrabertModel(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+
+        self.layers = nn.ModuleList(
+            [UltrabertLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = UltrabertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = nn.ModuleDict()
+        self.gradient_checkpointing = False
+        for base in set(config.rope_theta):
+            name = str(base).replace(".", "_")
+            self.rotary_emb[name] = RotaryEmbedding(base, config.head_dim)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutput:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("must specify exactly one of input_ids or inputs_embeds")
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if position_ids is None:
+            position_ids = torch.arange(inputs_embeds.size(1))
+            position_ids = position_ids.to(inputs_embeds.device).unsqueeze(0)
+        if (
+            attention_mask is None
+            and self.config._attn_implementation != "flash_attention_2"
+        ):
+            attention_mask = inputs_embeds.new_ones(inputs_embeds.shape[:2])
+        varlen_kwargs = {
+            k: kwargs[k]
+            for k in ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k")
+            if kwargs.get(k) is not None
+        }
+
+        # Create attention masks for different local window sizes and store by window.
+        attention_mask_dict = {}
+        for local_window in list(set(self.config.attention_windows)) + [None]:
+            if self.config._attn_implementation == "flash_attention_2":
+                # For FlashAttention 2, rely on the `sliding_window` keyword instead of
+                # constructing masks manually.
+                attention_mask_dict[local_window] = attention_mask
+            else:
+                window_mask = create_local_attention_mask(
+                    seq_len=attention_mask.size(1),
+                    sliding_window=local_window,
+                    device=attention_mask.device,
+                )
+                window_mask = window_mask & attention_mask.bool().unsqueeze(1)
+                attention_mask_dict[local_window] = (
+                    torch.where(window_mask, 0.0, torch.finfo(inputs_embeds.dtype).min)
+                    .to(inputs_embeds.dtype)
+                    .unsqueeze(1)
+                )
+
+        # Create position embeddings to be shared across the decoder layers.
+        hidden_states, position_embeddings = inputs_embeds, {}
+        for base in set(self.config.rope_theta):
+            rotary_emb = self.rotary_emb[str(base).replace(".", "_")]
+            position_embeddings[base] = rotary_emb(hidden_states, position_ids)
+
+        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            local_window = self.config.attention_windows[i]
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask_dict[local_window],
+                position_ids=position_ids,
+                position_embeddings=position_embeddings[self.config.rope_theta[i]],
+                sliding_window=(
+                    local_window + 1
+                    if local_window is not None
+                    and self.config._attn_implementation == "flash_attention_2"
+                    else local_window
+                ),
+                **varlen_kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class UltrabertPredictionHead(nn.Module):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__()
+        if (dim := config.prediction_head_size) is not None:
+            self.linear1 = nn.Linear(config.hidden_size, dim, bias=False)
+            self.linear2 = nn.Linear(dim, config.hidden_size, bias=False)
+            self.act_fn = ACT2FN["gelu_pytorch_tanh"]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "linear1"):
+            return self.act_fn(self.linear2(self.act_fn(self.linear1(x))))
+        return x
+
+
+@auto_docstring
+class UltrabertForMaskedLM(UltrabertPreTrainedModel):
+    _tied_weights_keys = {"cls.weight": "ultrabert.embed_tokens.weight"}
+
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.ultrabert = UltrabertModel(config)
+        self.head = UltrabertPredictionHead(config)
+        self.cls = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.cls
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls = new_embeddings
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor] | MaskedLMOutput:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
+            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+        """
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        prediction_scores = self.cls(self.head(outputs[0]))
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
+            )
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    Ultrabert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
+    output) e.g. for GLUE tasks.
+    """
+)
+class UltrabertForSequenceClassification(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.ultrabert = UltrabertModel(config)
+        classifier_dropout = (
+            config.classifier_dropout
+            if config.classifier_dropout is not None
+            else config.hidden_dropout_prob
+        )
+        self.dropout = nn.Dropout(config.classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        logits = self.classifier(self.dropout(outputs[0][:, 0]))
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring
+class UltrabertForQuestionAnswering(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.ultrabert = UltrabertModel(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        start_positions: torch.Tensor | None = None,
+        end_positions: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor] | QuestionAnsweringModelOutput:
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        logits = self.qa_outputs(outputs[0])
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring
+class UltrabertForTokenClassification(UltrabertPreTrainedModel):
+    def __init__(self, config: UltrabertConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.ultrabert = UltrabertModel(config)
+        classifier_dropout = (
+            config.classifier_dropout
+            if config.classifier_dropout is not None
+            else config.hidden_dropout_prob
+        )
+        self.dropout = nn.Dropout(config.classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor] | TokenClassifierOutput:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+        """
+        outputs = self.ultrabert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

model.safetensors

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

rl_agent_config.json

@@ -0,0 +1,54 @@
+{
+  "head_layers": 2,
+  "max_prefixes": 6,
+  "act_costs": {
+    "escalate": 0.5
+  },
+  "cost_wrong_act": 3.0,
+  "encoder": "dmis-lab/UltraBERT-150M",
+  "encoder_revision": "545ac9d9c7aece444f014ada5a1e9818132b4d35",
+  "max_len": 8192,
+  "head_max_len": 192,
+  "amp_dtype": "bf16",
+  "model_name": "caldec-ultra-150m",
+  "temperature": [
+    1.0,
+    1.0,
+    1.0
+  ],
+  "temperature_by_options": {
+    "choice:3-5": 1.3067494630813599,
+    "noul:2": 1.1769542694091797,
+    "score:3-5": 1.1716774702072144,
+    "choice:2": 1.1175421476364136
+  },
+  "training": {
+    "fine_tuned_from": "dmis-lab/UltraBERT-150M@545ac9d9c7aece444f014ada5a1e9818132b4d35 with a freshly initialized decision head",
+    "base_revision": "545ac9d9c7aece444f014ada5a1e9818132b4d35",
+    "epochs_completed": 6,
+    "selected_epoch": 6,
+    "updates": 3353,
+    "lr_encoder": 1e-05,
+    "lr_head": 0.0001,
+    "schedule": "linear",
+    "decisions": 9001,
+    "seconds": 2066,
+    "legacy_drop_tail": false,
+    "objective": "soft cross-entropy",
+    "with_public": true,
+    "balance_sites": true,
+    "seed": 20260924,
+    "packed": true,
+    "max_len": 8192,
+    "max_batch_tokens": 16384,
+    "extra_train": [
+      "general mix (scripts/general_data.py: MNLI, BoolQ)",
+      "kgrozdanovski/assistant-decisions long/train.jsonl"
+    ],
+    "extra_val": [
+      "general mix validation (scripts/general_data.py)",
+      "kgrozdanovski/assistant-decisions long/validation.jsonl"
+    ]
+  },
+  "packed": true
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_0": "<|endoftext|>",
+  "extra_1": "<|im_start|>",
+  "extra_10": "<issue_closed>",
+  "extra_11": "<jupyter_start>",
+  "extra_12": "<jupyter_text>",
+  "extra_13": "<jupyter_code>",
+  "extra_14": "<jupyter_output>",
+  "extra_15": "<|sep|>",
+  "extra_16": "<|mask|>",
+  "extra_2": "<|im_end|>",
+  "extra_3": "<repo_name>",
+  "extra_4": "<reponame>",
+  "extra_5": "<file_sep>",
+  "extra_6": "<filename>",
+  "extra_7": "<gh_stars>",
+  "extra_8": "<issue_start>",
+  "extra_9": "<issue_comment>",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 32768,
+  "model_specific_special_tokens": {
+    "extra_0": "<|endoftext|>",
+    "extra_1": "<|im_start|>",
+    "extra_10": "<issue_closed>",
+    "extra_11": "<jupyter_start>",
+    "extra_12": "<jupyter_text>",
+    "extra_13": "<jupyter_code>",
+    "extra_14": "<jupyter_output>",
+    "extra_15": "<|sep|>",
+    "extra_16": "<|mask|>",
+    "extra_2": "<|im_end|>",
+    "extra_3": "<repo_name>",
+    "extra_4": "<reponame>",
+    "extra_5": "<file_sep>",
+    "extra_6": "<filename>",
+    "extra_7": "<gh_stars>",
+    "extra_8": "<issue_start>",
+    "extra_9": "<issue_comment>"
+  },
+  "pad_token": "<|endoftext|>",
+  "sep_token": "<|sep|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}