| import os, json |
| import torch, torch.nn as nn, torch.nn.functional as F |
| from transformers import AutoModel, AutoTokenizer, AutoConfig |
|
|
|
|
| class ArabicDecisionModel(nn.Module): |
| """jeb - Arabic typed decision model. |
| |
| Each option is written into the prompt next to its own [MASK] marker; the model |
| scores at that position and softmaxes across the options of that question. The |
| answer space is therefore data, not architecture - new schemas need no retraining. |
| """ |
|
|
| def __init__(self, base='UBC-NLP/MARBERTv2', head_layers=2, local_dir=None): |
| super().__init__() |
| if local_dir: |
| self.tok = AutoTokenizer.from_pretrained(os.path.join(local_dir, 'tokenizer')) |
| cfg = AutoConfig.from_pretrained(os.path.join(local_dir, 'encoder')) |
| self.enc = AutoModel.from_config(cfg) |
| else: |
| self.tok = AutoTokenizer.from_pretrained(base) |
| self.enc = AutoModel.from_pretrained(base) |
| H = self.enc.config.hidden_size |
| layer = nn.TransformerEncoderLayer(d_model=H, nhead=12, dim_feedforward=H * 4, |
| batch_first=True, activation='gelu', dropout=0.1) |
| self.head = nn.TransformerEncoder(layer, num_layers=head_layers) |
| self.scorer = nn.Sequential(nn.Linear(H, H), nn.GELU(), nn.LayerNorm(H), nn.Linear(H, 1)) |
| self.mask_id = self.tok.mask_token_id |
|
|
| def build(self, state, question, max_len=256): |
| """Serialize state + question + options, each option preceded by [MASK].""" |
| s = ' | '.join(f'{k}: {v}' for k, v in state.items()) if isinstance(state, dict) else str(state) |
| crit = question['criteria'] |
| opts = list(crit.keys()) if isinstance(crit, dict) else list(crit) |
| parts = [self.tok.cls_token, s, self.tok.sep_token, question['instructions']] |
| for o in opts: |
| desc = crit[o] if isinstance(crit, dict) else o |
| parts += [self.tok.mask_token, f'{o}: {desc}'] |
| return ' '.join(parts), opts |
|
|
| def forward(self, input_ids, attention_mask, mask_positions, option_counts): |
| h = self.enc(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state |
| h = self.head(h, src_key_padding_mask=~attention_mask.bool()) |
| return [self.scorer(h[b, mask_positions[b][:option_counts[b]]]).squeeze(-1) |
| for b in range(h.size(0))] |
|
|
| @torch.no_grad() |
| def predict(self, state, questions, max_len=256): |
| """questions: {id: {type, instructions, criteria}} -> {id: {answer, confidence, probabilities}}""" |
| out = {} |
| for qid, q in questions.items(): |
| text, opts = self.build(state, q, max_len) |
| enc = self.tok([text], return_tensors='pt', truncation=True, max_length=max_len) |
| dev = next(self.parameters()).device |
| enc = {k: v.to(dev) for k, v in enc.items()} |
| pos = (enc['input_ids'][0] == self.mask_id).nonzero(as_tuple=True)[0][:len(opts)] |
| if len(pos) < len(opts): |
| out[qid] = {'error': 'options truncated - raise max_len or use fewer options'} |
| continue |
| logits = self(enc['input_ids'], enc['attention_mask'], [pos], [len(opts)])[0] |
| p = F.softmax(logits.float(), -1) |
| n, peak = len(opts), float(p.max()) |
| conf = max(0.0, min(1.0, (n * peak - 1) / (n - 1))) if n > 1 else 1.0 |
| out[qid] = { |
| 'answer': opts[int(p.argmax())], |
| 'confidence': round(conf, 4), |
| 'probabilities': {o: round(float(v), 4) for o, v in zip(opts, p)}, |
| } |
| return out |
|
|
|
|
| def load(path='.', device='cpu'): |
| """Load jeb from a local clone or snapshot_download() of IJyad/jeb.""" |
| from safetensors.torch import load_file |
| m = ArabicDecisionModel(local_dir=path) |
| m.load_state_dict(load_file(os.path.join(path, 'model.safetensors'))) |
| return m.to(device).eval() |
|
|
