File size: 3,935 Bytes
ad5ba3c 3a549ce ad5ba3c 3a549ce ad5ba3c 3a549ce ad5ba3c 3a549ce ad5ba3c 3a549ce ad5ba3c 3a549ce ad5ba3c 3a549ce ad5ba3c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | 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()
|