README.md

---
language:
- th
license: mit
base_model: jhu-clsp/mmBERT-base
tags:
- system-1
- system one
- jev
- typed-decisions
- modernbert
- calibration
- customer-support
- classification
- business
datasets:
- yothinS/thai-customer-support
pipeline_tag: text-classification
---

# Thai Jev: System 1 Decision Model (mmBERT)

An implementation of the **Jev System 1** decision architecture built on top of [`jhu-clsp/mmBERT-base`](https://huggingface.co/jhu-clsp/mmBERT-base) and fine-tuned on [`yothinS/thai-customer-support`](https://huggingface.co/datasets/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:

```text
[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

```python
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 = "yothinS/Aaron"
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

```json
{
  "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.