yothinS/Aaron

Thai Jev: System 1 Decision Model (mmBERT)

An implementation of the Jev System 1 decision architecture built on top of jhu-clsp/mmBERT-base and fine-tuned on yothinS/thai-customer-support.

Unlike standard classifiers tied to fixed, rigid output heads, this model evaluates dynamic runtime schemas in a single non-autoregressive forward pass. It supports three typed decision primitives:

• Choice: Categorical routing over arbitrary dynamic candidate strings with calibrated softmax distributions.
• Noul: Calibrated binary boolean verification ([0.0, 1.0]) for arbitrary hypotheses.
• Score: Continuous scalar evaluation ([0.0, 1.0]) guided by custom natural-language rubrics.

---

Model Architecture & Training

The architecture couples a shared multilingual ModernBERT transformer backbone with three specialized projection heads operating on mean-pooled token representations:

[State + Candidate / Rubric Prompt]
│
┌────────▼────────┐
│   mmBERT-base   │ (ModernBERT encoder)
└────────┬────────┘
         │ Mean Pooling
┌────────▼────────┐
│  d_model (768)  │
└────────┬────────┘
┌─────────┼─────────┐
▼         ▼         ▼
Choice   Noul      Score
Head     Head      Head
(GELU+MLP)(GELU+MLP)(GELU+Sigmoid)

Calibrated Multi-Task Loss

To prevent overconfidence and provide well-calibrated probabilities for downstream routing agents, training combines classification losses with a Brier score penalty:

• Brier Weight: 0.5
• Tail Reservation Strategy: A fixed reserve buffer (TAIL_RESERVE = 72) ensures ticket contexts are clipped before prompts and candidate options are appended, preventing candidate truncation.

---

Training Hyperparameters

Parameter	Setting
Base Model	jhu-clsp/mmBERT-base
Dataset	yothinS/thai-customer-support
Max Sequence Length	192 tokens
Epochs	3
Optimizer	AdamW ($\text{lr} = 2\times 10^{-5}$, $\text{weight decay} = 0.01$)
Learning Rate Schedule	Linear warmup with decay (Warmup ratio: 0.06)
Gradient Accumulation	8 steps
Precision	Mixed Precision (AMP torch.autocast)
Gradient Clipping	Max norm 1.0

---

Input Formats

Each decision type formats the input string prior to tokenization:

• Choice candidate format:
  สถานะ: {state} | คำถาม: {question} | ตัวเลือก: {option}
• Noul hypothesis format:
  สถานะ: {state} | สมมติฐาน: {hypothesis}
• Score rubric format:
  สถานะ: {state} | เกณฑ์ประเมิน: {rubric}

---

Quickstart & Inference

import os
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

# 1. Define Model Class
class aaronSystemOne(nn.Module):
    def __init__(self, backbone_id="jhu-clsp/mmBERT-base"):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(backbone_id)
        d_model = self.encoder.config.hidden_size

        def make_head(sigmoid=False):
            layers = [
                nn.Linear(d_model, d_model // 2),
                nn.GELU(),
                nn.Dropout(0.1),
                nn.Linear(d_model // 2, 1)
            ]
            if sigmoid:
                layers.append(nn.Sigmoid())
            return nn.Sequential(*layers)

        self.choice_head = make_head()
        self.noul_head = make_head()
        self.score_head = make_head(sigmoid=True)

    def extract_embedding(self, input_ids, attention_mask):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        mask = attention_mask.unsqueeze(-1).float()
        return (out.last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)

    def score_choices(self, input_ids, attention_mask):
        return self.choice_head(self.extract_embedding(input_ids, attention_mask)).squeeze(-1)

    def score_noul(self, input_ids, attention_mask):
        return self.noul_head(self.extract_embedding(input_ids, attention_mask)).squeeze(-1)

    def score_metric(self, input_ids, attention_mask):
        return self.score_head(self.extract_embedding(input_ids, attention_mask)).squeeze(-1)

# 2. Load Pipeline and Engine
# Note: Ensure modeling_jev.py is in your working directory
from modeling_jev import aaronEngine, Choice, Noul, Score

MODEL_ID = "YOUR_USERNAME/thai-jev-system-one-mmbert"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = aaronSystemOne(backbone_id=MODEL_ID)

# Support both local folder and remote Hugging Face repository
weights_path = f"{MODEL_ID}/model.safetensors"
if not os.path.exists(weights_path):
    weights_path = hf_hub_download(repo_id=MODEL_ID, filename="model.safetensors")

model.load_state_dict(load_file(weights_path))
model.to(device)
model.eval()

client = aaronEngine(model, tokenizer, device=device, max_length=192)
client.warmup()

# 3. Evaluate an Unseen Ticket
state = "ลูกค้าโวยวายหนักมาก บอกว่าบัตรเครดิตถูกชาร์จซ้ำซ้อน 3 รอบและขู่จะฟ้อง สคบ. ถ้าไม่ได้เงินคืนภายในวันนี้"

decision = client.judge(state, {
    "route_department": Choice(
        question="คำร้องนี้ควรส่งต่อไปยังทีมงานใด?",
        options=["ทีมกฎหมายและข้อพิพาท", "ทีมการเงินและบิลลิ่ง", "ทีมวิศวกรรมไอที"]
    ),
    "legal_threat": Noul("ลูกค้ามีการข่มขู่ว่าจะดำเนินการทางกฎหมายใช่หรือไม่"),
    "churn_risk": Score("ประเมินระดับความเสี่ยงที่ลูกค้าจะเลิกใช้บริการ (Churn Risk)")
})

print(decision)

Example Output

{
  "route_department": {
    "type": "Choice",
    "value": "ทีมการเงินและบิลลิ่ง",
    "probability": 0.8842,
    "confidence": 0.7931,
    "distribution": {
      "ทีมการเงินและบิลลิ่ง": 0.8842,
      "ทีมกฎหมายและข้อพิพาท": 0.0911,
      "ทีมวิศวกรรมไอที": 0.0247
    }
  },
  "legal_threat": {
    "type": "Noul",
    "value": true,
    "probability": 0.9412
  },
  "churn_risk": {
    "type": "Score",
    "value": 0.8625
  }
}

---

Intended Uses & Limitations

Best Use Cases

• Triage & Intent Routing: Fast dispatching of incoming Thai customer tickets to target departments.
• Guardrails & Escalation Checks: Instant binary verification of high-risk conditions (e.g., fraud claims, legal actions, VIP accounts).
• System 1 Filtering: Operating as a high-throughput, low-latency filter in front of slower autoregressive LLMs (System 2).

Limitations

• Maximum Context Length: The default sequence length is capped at 192 tokens. Extremely long conversations will have older turns truncated.
• Domain Specificity: The fine-tuning target represents customer support operations. Performance on generic open-domain tasks may be limited without further fine-tuning.
• Scoring Subjectivity: Score predictions depend directly on how clearly the rubric text reflects patterns represented during training.