Laya-HE v6 (seed 1)

README.md

@@ -0,0 +1,213 @@
+---
+license: cc-by-nc-sa-4.0
+language:
+- he
+- en
+library_name: laya
+pipeline_tag: zero-shot-classification
+base_model: dicta-il/neodictabert-bilingual
+tags:
+- laya
+- hebrew
+- decision-model
+- calibrated
+datasets:
+- LocalLLaMA/typed-decisions
+- clinc/clinc_oos
+- legacy-datasets/banking77
+- fancyzhx/ag_news
+- fancyzhx/dbpedia_14
+- community-datasets/yahoo_answers_topics
+- facebook/anli
+- Yelp/yelp_review_full
+- HebArabNlpProject/HebNLI
+- Etelis/HeQ_v1
+- HebArabNlpProject/HebrewSentiment
+- Tobi-Bueck/customer-support-tickets
+- nyu-mll/glue
+- google/civil_comments
+- google-research-datasets/go_emotions
+- ucberkeley-dlab/measuring-hate-speech
+- tasksource/tasksource-instruct-v0
+- wikimedia/wikipedia
+- HuggingFaceFW/fineweb-2
+- google/boolq
+- ehovy/race
+---
+
+# Laya-HE v6: a Hebrew–English decision model
+
+A [Laya](https://github.com/NandhaKishorM/laya)-style decision model for Hebrew and English. You give it a **state** (the
+text or fields to judge) and **questions**: a choice between options, a score on a scale, or a yes/no claim. In one
+forward pass it returns a **calibrated probability for every answer**. It is a fast classifier that you configure at
+call time. It is not a chatbot and it does not generate text.
+
+- **Encoder:** [`dicta-il/neodictabert-bilingual`](https://huggingface.co/dicta-il/neodictabert-bilingual) (NeoBERT, 28
+  layers, Hebrew + English)
+- **Head:** Laya's `DecisionModel` head, with 2 transformer layers and a scorer over the option markers
+- **Size:** 378M parameters, about 120–230 ms per call on an Apple M1 CPU
+- **Input:** up to 1,024 tokens of state, and up to 256 tokens of options
+- **Training:** Laya's RLCD objective (proper-scoring-rule rewards plus soft cross-entropy), with a temperature for each
+  question type fitted on held-out human- or rule-labeled items
+- **Use:** non-commercial only (see [License](#license))
+
+## Usage
+
+This checkpoint needs Laya 0.3.7 (commit `010bace`) with `neobert.patch`. The patch loads NeoBERT's remote code,
+recomputes its rotary tables (without that, every output is NaN under transformers 5) and keeps the encoder in fp32.
+
+```bash
+git clone https://github.com/NandhaKishorM/laya && cd laya && git checkout 010bace
+git apply /path/to/neobert.patch          # from this repository
+pip install -e .                          # torch, transformers 5.x
+```
+
+```python
+import laya
+
+agent = laya.load("RoeiG/laya-hebrew", device="cpu")   # or "cuda"
+out = agent.predict(
+    {"message": "האפליקציה קורסת כשאני פותח את המצלמה"},
+    {
+        "team": {"type": "choice", "instructions": "איזה צוות צריך לטפל בהודעה?",
+                 "criteria": {"billing": "תשלומים והחזרים", "tech": "באגים וקריסות", "shipping": "משלוחים"}},
+        "upset": {"type": "noul", "instructions": "הלקוח כועס."},
+    },
+)
+# out["answers"]["team"]  -> {"choice": "tech", "probabilities": {"billing": 0.0014, "tech": 0.9973, "shipping": 0.0013}, ...}
+# out["answers"]["upset"] -> {"noul": 0.0551, ...}   (P(true))
+```
+
+Question types:
+- **`choice`:** `criteria` maps each option to a description.
+- **`score`:** `criteria` is a list of levels from lowest to highest. The answer includes the expected level.
+- **`noul`:** yes/no. The answer is P(true) for the statement in `instructions`.
+
+## How to get good answers
+
+1. **Compute numbers, dates, units and relations in code.** Pass the result as a field, such as
+   `"age_ok": "הגיל עומד בתנאי"` or `"relation": "The sender is the receiver's direct manager"`. The model does not do
+   arithmetic reliably (see Limitations).
+2. **Prefer the claim form for yes/no.** "הלקוח כועס." discriminates better than "האם הלקוח כועס?" (gap 0.67 against
+   0.38 on he_bench). The question form works in v6, but it is weaker.
+3. **Describe every option in a line.** Bare labels or codes route much worse than labels with a one-line description.
+   Name what each option owns, not only its keywords.
+4. **Read the probabilities, not only the top answer.** A top answer below about 0.6 means the model is unsure.
+5. **Ignore `act_probability`.** It comes from a head that no Hebrew checkpoint trained.
+
+## Evaluation
+
+Every evaluation set below was held out of training. he_bench v1 has 8 Hebrew tasks, each asked in 3 phrasings; the
+set is frozen by sha256. Brier and ECE are better when lower. "laya-multilingual" is Laya's published multilingual
+checkpoint. v5 is the previous release of this project.
+
+| Metric | laya-multilingual | v5 | **v6** |
+|---|---|---|---|
+| MASSIVE he, 20 intents (500) | 0.352 | 0.816 | **0.806** |
+| MASSIVE he, 4 intents | 0.710 | 0.928 | **0.938** |
+| MASSIVE en, 20 intents | 0.652 | 0.816 | **0.816** |
+| he_bench accuracy / Brier / ECE (4,290) | 0.439 / 0.701 / 0.266 | 0.527 / 0.538 / 0.119 | **0.595 / 0.499 / 0.125** |
+| Belebele-he reading comprehension (900; chance 0.25) | | 0.468 | **0.767** |
+| SIB-200-he topic | | 0.808 | **0.801** |
+| Yes/no as a question: P(yes \| true) − P(yes \| false) | | 0.047 | **0.379** |
+| Yes/no as a claim: same gap (390) | | 0.692 | **0.674** |
+| Hebrew BoolQ, held out (875) | | 0.611 | **0.838** |
+| Rule-direction probe (208) | | 0.798 | **0.832** |
+| Held-out soft-label set, Brier (900) | | 0.230 | **0.231** |
+
+he_bench accuracy by task (v6):
+
+| Task | Accuracy | Chance |
+|---|---|---|
+| relevance (question form) | 0.785 | 0.50 |
+| qa_verify (question form) | 0.704 | 0.50 |
+| copa | 0.687 | 0.50 |
+| sentiment | 0.633 | 0.33 |
+| winograd | 0.607 | 0.50 |
+| hellaswag | 0.447 | 0.25 |
+| tone arousal | 0.389 | 5 levels |
+| tone valence | 0.328 | 5 levels |
+
+**Run-to-run noise.** A second training seed, with the same data and settings, differed by these amounts:
+- he_bench overall: 0.2 points
+- Belebele and MASSIVE: 0.6–1.0 points
+- single he_bench tasks: 1–3 points
+- each half of the rule probe: 5–7 points
+
+Differences smaller than these are noise. This checkpoint is seed 1, which was fixed as the release before training.
+
+## Limitations
+
+- **Reasoning is the weak spot.**
+  - Hellaswag, winograd and copa are well above chance but far from solved.
+  - Multi-step inferences (e.g. "A is taller than B, B is taller than C: who is shortest?") often fail.
+- **Numeric, date and unit rules are unreliable, and often confidently wrong.**
+  - Examples: 2.5 hours against a 2-hour limit, a purchase 19 days ago against a 14-day window, or age 17 against an
+    English "18 and up" rule. These got P(true) of 0.93–0.98.
+  - Compute them in code.
+- **Irony and sarcasm are read literally.** "וואו, שירות מדהים… ניתקו לי 👏" is scored as positive, at 0.97.
+- **Routing leans on keywords.**
+  - In a small hand-written check of tech tickets, anything that mentioned "דיפלוי" (deploy) was pulled towards DevOps.
+    That happened even when the cause was a code bug, a network path, an expired certificate or a locked account.
+  - Strong keywords in the text can outweigh the option descriptions.
+- **Claim-form relevance is slightly below v5:** he_bench 0.78 against 0.80, and BEIR-he 0.78 against 0.82.
+- **Scales:** sentiment and tone are about 0.45 accuracy, and the tone probabilities are overconfident (valence Brier
+  0.53). The lowest level of a scale is under-used.
+- **Calculated fields steer less than in v5.** On 8 test emails, a "direct manager" sender field raised importance by
+  0.11 of a level, against 0.22 in v5, and a "mailing list" field lowered it in only 2 of 8.
+- **Calibration does not catch everything.** The failures above are often high-confidence, so a confidence threshold
+  will not filter them out.
+
+## Training data
+
+The training start was an earlier checkpoint of this project (v3), trained on the same public data and on the same
+encoder. The run was one epoch over 396,508 items (6,196 updates, 1.4 A100 hours). The learning rates were 5e-6 for the
+encoder and 1e-4 for the head. Every case was converted to Laya's format, with a random subset of options, paraphrased
+instructions and varied field names.
+
+- **Public labeled data**, English unless marked:
+  - typed-decisions
+  - CLINC, Banking77, AG News, DBpedia, Yahoo Answers
+  - ANLI, Yelp, GLUE STS-B, Civil Comments, customer-support tickets
+  - Hebrew NLI (HebNLI), Hebrew QA (HeQ) and Hebrew sentiment
+- **Soft labels from annotator disagreement:**
+  - GoEmotions and Measuring Hate Speech
+  - SNLI and MultiNLI votes, used as P(true)
+  - plus breadth from tasksource-instruct
+- **Native Hebrew:** Hebrew Wikipedia topics and facts, labeled from Wikidata.
+- **Relevance and spam:** BEIR-he relevance (biunlp's Hebrew translation of BEIR) and UCI SMS spam.
+- **Machine-translated:** part of the English data was translated to Hebrew with NLLB-200-distilled-600M: 60% of the
+  soft-label cases, and half of the NLI votes and the spam. Only the state text was translated; the questions and
+  options stayed as they were.
+- **Generated with code-computed labels:** rule and unit checks.
+- **New in v6:**
+  - RACE (30,000 questions) and BoolQ, in English with human labels.
+  - About 58,000 Hebrew items written and labeled by
+    [DictaLM-3.0-24B](https://huggingface.co/dicta-il) (Apache-2.0), on passages from Hebrew Wikipedia and FineWeb-2
+    Hebrew:
+    - yes/no questions about a passage, and answer checking
+    - query–passage relevance
+    - 4-option reading comprehension
+    - routing of invented business messages
+    - BoolQ translated to Hebrew, which keeps its human answers
+  - An item was kept only when the teacher's label agreed with the answer it was written for.
+  - Subjective scales were never teacher-labeled.
+  - The teacher-generated files are not published.
+
+Decontamination: no teacher text shares an 8-word run with any evaluation set. The evaluation instructions and the
+yes/no wordings used by he_bench were kept out of training. MASSIVE was never trained on.
+
+No private data and no personal data were used.
+
+## License
+
+**CC-BY-NC-SA-4.0: non-commercial use only.** Several training sources are non-commercial or research-only: ANLI, Yelp,
+AG News, Yahoo Answers, RACE, and the NLLB translation model (CC-BY-NC-4.0). Others are ShareAlike: Wikipedia, BEIR-he,
+SNLI and BoolQ. The encoder is by Dicta (`neodictabert-bilingual`, CC-BY-4.0). The Laya architecture, training method
+and runtime are by Laya's authors (Apache-2.0).
+
+## Acknowledgements
+
+- Dicta, for NeoDictaBERT-bilingual and DictaLM 3.0
+- Laya's authors, for the architecture, the RLCD training method and the runtime
+- The creators of every dataset listed above

encoder/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "NeoBERTLMHead"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_neobert.NeoBERTConfig",
+    "AutoModel": "modeling_neobert.NeoBERT",
+    "AutoModelForMaskedLM": "modeling_neobert.NeoBERTLMHead",
+    "AutoModelForQuestionAnswering": "modeling_neobert.NeoBERTForQuestionAnswering",
+    "AutoModelForSequenceClassification": "modeling_neobert.NeoBERTForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_neobert.NeoBERTForTokenClassification"
+  },
+  "decoder_init_range": 0.02,
+  "dim_head": 64,
+  "dtype": "bfloat16",
+  "embedding_init_range": 0.02,
+  "encoder_init_range": 0.02,
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "max_length": 4096,
+  "model_type": "neobert",
+  "norm_eps": 1e-06,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "pad_token_id": 3,
+  "transformers_version": "5.16.1",
+  "vocab_size": 128000
+}

encoder/modeling_neobert.py

@@ -0,0 +1,696 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn.functional import scaled_dot_product_attention
+
+from typing import Optional, Tuple, Union
+import numpy as np
+
+try:
+    from xformers.ops import SwiGLU
+except:
+    class SwiGLU(nn.Module):
+        """
+        A Module that mimicks the call to :attr:`xformers.ops.swiglu`,
+        and holds the weights for the 3 linear layers
+        """
+        def __init__(
+            self,
+            in_features: int,
+            hidden_features: int,
+            out_features: Optional[int] = None,
+            bias: bool = True,
+            *,
+            _pack_weights: bool = True,
+        ) -> None:
+            """Create a SwiGLU module
+
+            Args:
+                in_features (int): Number of features of the input
+                hidden_features (int): Number of hidden features
+                out_features (Optional[int], optional): Number of features of the input. Defaults to None.
+                bias (bool, optional): Whether linear layers also include a bias. Defaults to True.
+            """
+            super().__init__()
+            out_features = out_features or in_features
+            hidden_features = hidden_features or in_features
+
+            self.w12: Optional[nn.Linear]
+            if _pack_weights:
+                self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+            else:
+                self.w12 = None
+                self.w1 = nn.Linear(in_features, hidden_features, bias=bias)
+                self.w2 = nn.Linear(in_features, hidden_features, bias=bias)
+            self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+            self.hidden_features = hidden_features
+            self.out_features = out_features
+            self.in_features = in_features
+            self.op: Optional[SwiGLUOp] = None
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Computes :attr:`swiglu` with the module's weights
+
+            Args:
+                x (torch.Tensor): A Tensor of shape ``[..., in_features]``
+
+            Returns:
+                torch.Tensor: A Tensor of shape ``[..., out_features]``
+            """
+            if self.w12 is not None:
+                gate, x = self.w12(x).chunk(2, dim=-1)
+                hidden = F.silu(gate) * x
+            else:
+                x1 = self.w1(x)
+                x2 = self.w2(x)
+                hidden = F.silu(x1) * x2
+
+            return self.w3(hidden)
+
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+from transformers import (
+    PreTrainedModel,
+    PretrainedConfig,
+    DataCollatorForLanguageModeling,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+    QuestionAnsweringModelOutput
+)
+
+import torch
+from typing import Tuple
+
+def precompute_freqs(dim: int, end: int, theta: float = 10000.0, *, device=None, dtype=torch.float32):
+    """
+    Returns (cos, sin) tensors of shape [end, dim//2], no complex dtype.
+    """
+    h = dim // 2
+    idx = torch.arange(0, h, device=device, dtype=dtype)
+    inv_freq = 1.0 / (theta ** ((2.0 * idx) / dim))
+    t = torch.arange(end, device=device, dtype=dtype)
+    angles = torch.outer(t, inv_freq)                         # [L, h]
+    return angles.cos(), angles.sin()   # ([L, h], [L, h])
+
+def reshape_for_broadcast(freqs: torch.Tensor, x: torch.Tensor):
+    # freqs: [L, h]; x: [B, L, H, h] for the half-dim tensors
+    assert freqs.shape == (x.shape[1], x.shape[-1]), (freqs.shape, x.shape)
+    return freqs[None, :, None, :]                            # [1, L, 1, h]
+
+# Rotary embedding without complex numbers (megatron-core pairing: first half with second half)
+def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]):
+    # x*: [B, L, H, D]; freqs = (cos[L,h], sin[L,h])
+    D = xq.shape[-1]
+    h = D // 2
+    xq1, xq2 = xq[..., :h], xq[..., h:]
+    xk1, xk2 = xk[..., :h], xk[..., h:]
+
+    cos, sin = freqs
+    cos = reshape_for_broadcast(cos.type_as(xq1), xq1)        # [1, L, 1, h]
+    sin = reshape_for_broadcast(sin.type_as(xq1), xq1)        # [1, L, 1, h]
+
+    q1 = xq1 * cos - xq2 * sin
+    q2 = xq1 * sin + xq2 * cos
+    k1 = xk1 * cos - xk2 * sin
+    k2 = xk1 * sin + xk2 * cos
+
+    return torch.cat([q1, q2], dim=-1), torch.cat([k1, k2], dim=-1)
+
+class NeoBERTEagerRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        NeoBERTEagerRMSNorm is equivalent to nn.RMSNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        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 NeoBERTConfig(PretrainedConfig):
+    model_type = "neobert"
+
+    # All config parameters must have a default value.
+    def __init__(
+        self,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        embedding_init_range: float = 0.02,
+        encoder_init_range: float = 0.02,
+        norm_eps: float = 1e-06,
+        vocab_size: int = 30522,
+        pad_token_id: int = 0,
+        max_length: int = 1024,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError("Hidden size must be divisible by the number of heads.")
+        self.dim_head = hidden_size // num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.embedding_init_range = embedding_init_range
+        self.encoder_init_range = encoder_init_range
+        self.norm_eps = norm_eps
+        self.vocab_size = vocab_size
+        self.pad_token_id = pad_token_id
+        self.max_length = max_length
+        self.kwargs = kwargs
+
+
+class EncoderBlock(nn.Module):
+    """Transformer encoder block."""
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__()
+
+        self.config = config
+
+        # Attention
+        self.qkv = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size * 3, bias=False)
+        self.wo = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=False)
+
+        # Feedforward network
+        # Original NeoBERT:
+        # multiple_of = 8
+        # intermediate_size = int(2 * config.intermediate_size / 3)
+        # intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
+        intermediate_size = config.intermediate_size
+        self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=False)
+
+        # Layer norms
+        rms_norm_cls = nn.RMSNorm if hasattr(nn, 'RMSNorm') else NeoBERTEagerRMSNorm
+        self.attention_norm = rms_norm_cls(config.hidden_size, config.norm_eps)
+        self.ffn_norm = rms_norm_cls(config.hidden_size, config.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        output_attentions: bool,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        # Attention
+        attn_output, attn_weights = self._att_block(
+            self.attention_norm(x), attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens
+        )
+
+        # Residual
+        x = x + attn_output
+
+        # Feed-forward
+        x = x + self.ffn(self.ffn_norm(x))
+
+        return x, attn_weights
+
+    def _att_block(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        output_attentions: bool,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        batch_size, seq_len, _ = x.shape
+
+        xq, xk, xv = self.qkv(x).view(batch_size, seq_len, self.config.num_attention_heads, self.config.dim_head * 3).chunk(3, axis=-1)
+
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+
+        # Attn block
+        attn_weights = None
+
+        # Flash attention if the tensors are packed
+        if cu_seqlens is not None:
+            attn = flash_attn_varlen_func(
+                q=xq.squeeze(0),
+                k=xk.squeeze(0),
+                v=xv.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                dropout_p=0.0,
+                causal=False,
+            )
+        # Eager attention if attention weights are needed in the output (avoid using this unless needed - e.g., for onnx export)
+        elif output_attentions or self.config._attn_implementation == 'eager':
+            attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
+            if attention_mask is not None:
+               attn_weights = attn_weights * attention_mask
+            attn_weights = attn_weights.softmax(-1)
+            attn = attn_weights @ xv.permute(0, 2, 1, 3)
+            attn = attn.transpose(1, 2)
+        # Fall back to SDPA otherwise
+        else:
+            attn = scaled_dot_product_attention(
+                query=xq.transpose(1, 2),
+                key=xk.transpose(1, 2),
+                value=xv.transpose(1, 2),
+                attn_mask=attention_mask.bool(),
+                dropout_p=0,
+            ).transpose(1, 2)
+
+        return self.wo(attn.reshape(batch_size, seq_len, self.config.num_attention_heads * self.config.dim_head)), attn_weights
+
+
+class NeoBERTPreTrainedModel(PreTrainedModel):
+    config_class = NeoBERTConfig
+    base_model_prefix = "model"
+    _supports_cache_class = True
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.uniform_(-self.config.encoder_init_range, self.config.encoder_init_range)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)
+
+
+class NeoBERT(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+
+        # Ensures freqs_cis is moved to the same devices as the model. Non-persistent buffers are not saved in the state_dict.
+        cos, sin = precompute_freqs(config.hidden_size // config.num_attention_heads, config.max_length)
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+
+        self.transformer_encoder = nn.ModuleList()
+        for _ in range(config.num_hidden_layers):
+            self.transformer_encoder.append(EncoderBlock(config))
+
+        rms_norm_cls = nn.RMSNorm if hasattr(nn, 'RMSNorm') else NeoBERTEagerRMSNorm
+        self.layer_norm = rms_norm_cls(config.hidden_size, config.norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None, # kept in to not break compatibility with tokenizer(...), ignored
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+        # Initialize
+        hidden_states, attentions = [], []
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        attention_mask = attention_mask[:, None, None, :]
+            
+            # attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
+
+        # Checks to be done if inputs are packed sequences
+        if cu_seqlens is not None:
+            assert (
+                FLASH_ATTN_AVAILABLE
+            ), "Flash-attention is not available. Please ''pip install flash_attn'', or provide un-packed sequences."
+            assert not output_attentions, "Output attentions is not supported when sequences are packed."
+            assert max_seqlen is not None, "Missing max_seqlen. It must be provided when cu_seqlens are not None."
+            assert (input_ids if input_ids is not None else inputs_embeds).shape[
+                0
+            ] == 1, "Cumulative sequence lengths are provided but inputs are not packed."
+            assert (
+                input_ids if input_ids is not None else inputs_embeds
+            ).is_cuda, "Packing uses an implementation of flash-attention and is only supported on GPU."
+
+        # RoPE
+        if position_ids is not None:
+            freqs = (self.freqs_cos[position_ids], self.freqs_sin[position_ids])
+        else:
+            L = (input_ids if input_ids is not None else inputs_embeds).shape[1]
+            freqs = (self.freqs_cos[:L], self.freqs_sin[:L])
+
+        # Embedding
+        x = self.encoder(input_ids) if input_ids is not None else inputs_embeds
+
+        # Transformer encoder
+        for layer in self.transformer_encoder:
+            x, attn = layer(x, attention_mask, freqs, output_attentions, max_seqlen, cu_seqlens)
+            if output_hidden_states:
+                hidden_states.append(x)
+            if output_attentions:
+                attentions.append(attn)
+
+        # Final normalization layer
+        x = self.layer_norm(x)
+
+        # Return the output of the last hidden layer
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states if output_hidden_states else None,
+            attentions=attentions if output_attentions else None,
+        )
+
+
+class NeoBERTLMHead(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.model = NeoBERT(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None, # kept in to not break compatibility with tokenizer(...), ignored
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+
+        output = self.model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            max_seqlen=max_seqlen,
+            cu_seqlens=cu_seqlens,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        logits = self.decoder(output.last_hidden_state)
+
+        return MaskedLMOutput(
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+            logits=logits,
+        )
+
+
+class NeoBERTForTokenClassification(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.num_labels = getattr(config, "num_labels", 2)
+        self.classifier_dropout = getattr(config, "classifier_dropout", 0.1)
+        self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)
+
+        self.model = NeoBERT(config)
+
+        self.dense = nn.Linear(self.config.hidden_size, self.config.hidden_size)
+        self.dropout = nn.Dropout(self.classifier_dropout)
+        self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)
+
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None, # kept in to not break compatibility with tokenizer(...), ignored
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        output = self.model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            max_seqlen=max_seqlen,
+            cu_seqlens=cu_seqlens,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        hidden_states = output.last_hidden_state
+
+        x = self.dropout(hidden_states)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+
+        logits = self.classifier(x)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)[active_loss]
+                active_labels = labels.view(-1)[active_loss]
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            result = (logits,)
+            return ((loss,) + result) if loss is not None else result
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+        )
+
+
+class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.num_labels = getattr(config, "num_labels", 2)
+        self.classifier_dropout = getattr(config, "classifier_dropout", 0.1)
+        self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)
+
+        self.model = NeoBERT(config)
+
+        self.dense = nn.Linear(self.config.hidden_size, self.config.hidden_size)
+        self.dropout = nn.Dropout(self.classifier_dropout)
+        self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)
+
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None, # kept in to not break compatibility with tokenizer(...), ignored
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        output = self.model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            max_seqlen=max_seqlen,
+            cu_seqlens=cu_seqlens,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        hidden_states = output.last_hidden_state
+
+        x = hidden_states[:, 0, :]
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+
+        logits = self.classifier(x)
+
+        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)
+
+        if not return_dict:
+            result = (logits,)
+            return ((loss,) + result) if loss is not None else result
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+        )
+
+class NeoBERTForQuestionAnswering(NeoBERTPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.model = NeoBERT(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None, # kept in to not break compatibility with tokenizer(...), ignored
+        start_positions: Optional[torch.Tensor] = None,
+        end_positions: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        return_dict: Optional[bool] = None,
+    ) -> Union[tuple[torch.Tensor], QuestionAnsweringModelOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if output_attentions or output_hidden_states: return_dict = True
+
+        output = self.model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            max_seqlen=max_seqlen,
+            cu_seqlens=cu_seqlens,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True
+        )
+        hidden_states = output.last_hidden_state
+
+        logits = self.qa_outputs(hidden_states)
+        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 we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            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
+
+        if not return_dict:
+            output = (start_logits, end_logits)
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=output.hidden_states,
+            attentions=output.attentions,
+        )

model.safetensors

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

neobert.patch

@@ -0,0 +1,47 @@
+diff --git a/laya/agent.py b/laya/agent.py
+index 70163bf..b7b9a0b 100644
+--- a/laya/agent.py
++++ b/laya/agent.py
+@@ -186,7 +186,8 @@ class Agent:
+                 self.device = torch.device("cpu")
+ 
+         tok_dir = os.path.join(model_dir, "tokenizer")
+-        self.tok = AutoTokenizer.from_pretrained(tok_dir if os.path.exists(tok_dir) else self.cfg.get("encoder"))
++        self.tok = AutoTokenizer.from_pretrained(tok_dir if os.path.exists(tok_dir) else self.cfg.get("encoder"),
++                                                 trust_remote_code=bool(self.cfg.get("trust_remote_code", False)))
+ 
+         enc_dir = os.path.join(model_dir, "encoder")
+         # The checkpoint supplies every parameter; skip random/base-model weights.
+diff --git a/laya/common.py b/laya/common.py
+index 950c41d..6a633f7 100644
+--- a/laya/common.py
++++ b/laya/common.py
+@@ -139,11 +139,25 @@ class DecisionModel(nn.Module):
+ def build_model(cfg: Dict, encoder_dir: Optional[str] = None, pretrained: bool = True) -> DecisionModel:
+     from transformers import AutoConfig, AutoModel
+ 
++    # Encoders whose modeling code lives on the Hub (e.g. NeoBERT) must opt in via the config.
++    remote = bool(cfg.get("trust_remote_code", False))
+     if not pretrained or (encoder_dir and os.path.exists(encoder_dir)):
+-        ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"])
+-        enc = AutoModel.from_config(ecfg, attn_implementation="sdpa")
++        ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"], trust_remote_code=remote)
++        enc = AutoModel.from_config(ecfg, attn_implementation="sdpa", trust_remote_code=remote)
+     else:
+-        enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa")
++        enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa", trust_remote_code=remote)
++    # NeoBERT computes its RoPE tables as non-persistent buffers in __init__; transformers 5 builds models on
++    # the meta device, so they come back as uninitialised memory and every forward pass returns NaN.
++    if hasattr(enc, "freqs_cos"):
++        import sys
++
++        precompute = sys.modules[type(enc).__module__].precompute_freqs
++        c = enc.config
++        enc.freqs_cos, enc.freqs_sin = precompute(c.hidden_size // c.num_attention_heads, c.max_length)
++    # Some encoder configs pin a half-precision dtype (NeoBERT ships bfloat16), which transformers 5 honours. Keep
++    # master weights in float32: bf16 weights round most optimizer updates to zero, and CPU inference in bf16 is ~8x slower.
++    # Mixed precision is applied by autocast at train/inference time instead.
++    enc = enc.float()
+     return DecisionModel(enc, cfg.get("head_layers", 2), len(cfg.get("act_costs", {})) + 1)
+ 
+ 

rl_agent_config.json

@@ -0,0 +1,38 @@
+{
+  "encoder": "dicta-il/neodictabert-bilingual",
+  "head_layers": 2,
+  "max_len": 1024,
+  "head_max_len": 256,
+  "max_prefixes": 6,
+  "act_costs": {
+    "escalate": 0.5
+  },
+  "cost_wrong_act": 3.0,
+  "amp_dtype": "bf16",
+  "model_name": "laya-he-neodictabert-v6",
+  "temperature": [
+    1.0730161666870117,
+    1.1290489435195923,
+    1.1732172966003418
+  ],
+  "training": {
+    "init": "/content/drive/MyDrive/laya_he_v3/model.safetensors",
+    "seed": 1,
+    "items": 396508,
+    "updates": 6196,
+    "epochs_completed": 1,
+    "lr_encoder": 5e-06,
+    "lr_head": 0.0001,
+    "hours": 1.42,
+    "calibration_items": {
+      "0": 600,
+      "1": 600,
+      "2": 600
+    },
+    "avg_loss": 0.3189
+  },
+  "trust_remote_code": true,
+  "gradient_checkpointing": true,
+  "max_tokens_per_batch": 4096,
+  "fine_tuned": true
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "eos_token": "[SEP]",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}