| """Input layouts, heads and pooling for per-candidate scorers. Pure PyTorch, no SDK imports: |
| `openjev tod export --hf` copies this file into the exported model, so the SDK and the |
| exported model share one implementation. |
| |
| Two layouts read a decision (a query, and the candidates to score against it): |
| |
| pairs one sequence per candidate, tokenizer(query, candidate); the pooled state is the |
| mean over the whole pair. Works with any model and attention implementation. |
| packed one sequence per decision: [cls] query [sep] | cand_1 [sep] | cand_2 [sep] ... |
| The query attends to itself; each candidate span attends to the query and to |
| itself, never to another candidate (a block mask), and its positions restart |
| where the query ends, so a candidate's score does not change when other |
| candidates are added, removed or reordered. The query is encoded once per |
| decision, and each candidate is pooled over its own span. Needs a model that |
| takes `position_ids` and a 4D attention mask (checked by `packing_supported`). |
| |
| Heads map a pooled state to one logit: `linear`, or `mlp` (Linear-GELU-LayerNorm-Linear). |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| LAYOUTS = ('packed', 'pairs') |
| HEADS = ('mlp', 'linear') |
|
|
|
|
| def make_head(kind, hidden): |
| if kind == 'linear': |
| return torch.nn.Linear(hidden, 1) |
| if kind == 'mlp': |
| return torch.nn.Sequential(torch.nn.Linear(hidden, hidden), torch.nn.GELU(), torch.nn.LayerNorm(hidden), |
| torch.nn.Linear(hidden, 1)) |
| raise ValueError(f'head must be one of {HEADS}') |
|
|
|
|
| def special_ids(tokenizer): |
| """(cls, sep): the tokens that open the query and close each span (BOS/EOS for decoders).""" |
| cls = tokenizer.cls_token_id if tokenizer.cls_token_id is not None else tokenizer.bos_token_id |
| sep = tokenizer.sep_token_id if tokenizer.sep_token_id is not None else tokenizer.eos_token_id |
| return cls, sep |
|
|
|
|
| def pack(query_ids, candidate_ids, cls, sep, max_length): |
| """One packed decision: dict(ids, positions, spans). The query keeps at most half of |
| `max_length`; each candidate is cut so that query + candidate fits in `max_length`.""" |
| prefix = ([cls] if cls is not None else []) + list(query_ids)[:max_length // 2] + ([sep] if sep is not None else []) |
| ids, positions, spans = list(prefix), list(range(len(prefix))), [] |
| room = max(1, max_length - len(prefix) - (sep is not None)) |
| for c in candidate_ids: |
| span = list(c)[:room] + ([sep] if sep is not None else []) |
| start = len(ids) |
| ids += span |
| positions += range(len(prefix), len(prefix) + len(span)) |
| spans.append((start, len(ids))) |
| return dict(ids=ids, positions=positions, spans=spans, prefix=len(prefix)) |
|
|
|
|
| def collate_packed(items, pad_id): |
| """Tensors for a forward over packed decisions, plus `pool_index` (each token's |
| candidate number, or `count` for tokens that are not in a candidate span).""" |
| width = max(len(x['ids']) for x in items) |
| count = sum(len(x['spans']) for x in items) |
| ids = torch.full((len(items), width), pad_id, dtype=torch.long) |
| positions = torch.zeros((len(items), width), dtype=torch.long) |
| mask = torch.zeros((len(items), 1, width, width), dtype=torch.bool) |
| pool = torch.full((len(items), width), count, dtype=torch.long) |
| n = 0 |
| for b, x in enumerate(items): |
| length, p = len(x['ids']), x['prefix'] |
| ids[b, :length] = torch.tensor(x['ids']) |
| positions[b, :length] = torch.tensor(x['positions']) |
| mask[b, 0, :p, :p] = True |
| for s, e in x['spans']: |
| mask[b, 0, s:e, :p] = True |
| mask[b, 0, s:e, s:e] = True |
| pool[b, s:e] = n |
| n += 1 |
| idx = torch.arange(length, width) |
| mask[b, 0, idx, idx] = True |
| return dict(input_ids=ids, position_ids=positions, attention_mask=mask), dict(pool_index=pool, count=count) |
|
|
|
|
| def pooled_spans(hidden, pool_index, count): |
| """Mean of each candidate's span: (count, hidden) in float32.""" |
| flat = hidden.float().reshape(-1, hidden.shape[-1]) |
| index = pool_index.reshape(-1) |
| sums = torch.zeros((count + 1, flat.shape[1]), dtype=flat.dtype, device=flat.device).index_add_(0, index, flat) |
| counts = torch.zeros(count + 1, dtype=flat.dtype, device=flat.device).index_add_( |
| 0, index, torch.ones_like(index, dtype=flat.dtype)) |
| return sums[:count] / counts[:count].clamp(min=1).unsqueeze(1) |
|
|
|
|
| def pooled_mean(hidden, attention_mask): |
| mask = attention_mask.unsqueeze(-1).float() |
| return (hidden.float() * mask).sum(1) / mask.sum(1).clamp(min=1) |
|
|
|
|
| class Scorer(torch.nn.Module): |
| """trunk (a Hugging Face base model) + head: one logit per candidate.""" |
|
|
| def __init__(self, trunk, head): |
| super().__init__() |
| self.trunk, self.head = trunk, head |
|
|
| def forward(self, pool_index=None, count=None, **batch): |
| hidden = self.trunk(**batch).last_hidden_state |
| pooled = pooled_mean(hidden, batch['attention_mask']) if pool_index is None else \ |
| pooled_spans(hidden, pool_index, count) |
| return self.head(pooled.to(next(self.head.parameters()).dtype)).squeeze(-1).float() |
|
|
|
|
| def encode(tokenizer, layout, query, candidates, max_length, max_chars): |
| """Tokenize one decision for `layout`: a list of pair encodings, or one packed dict.""" |
| candidates = [c[:max_chars] for c in candidates] |
| if layout == 'pairs': |
| encoded = tokenizer([query] * len(candidates), candidates, truncation='longest_first', max_length=max_length) |
| return [{k: encoded[k][i] for k in encoded.keys()} for i in range(len(candidates))] |
| cls, sep = special_ids(tokenizer) |
| q = tokenizer(query, add_special_tokens=False)['input_ids'] |
| cs = tokenizer(candidates, add_special_tokens=False)['input_ids'] if candidates else [] |
| return pack(q, cs, cls, sep, max_length) |
|
|
|
|
| def size(layout, item): |
| """(sequences, longest) of one encoded decision, for budgeting padded tokens.""" |
| if layout == 'pairs': |
| return len(item), max(len(r['input_ids']) for r in item) |
| return 1, len(item['ids']) |
|
|
|
|
| def useful(layout, item): |
| return sum(len(r['input_ids']) for r in item) if layout == 'pairs' else len(item['ids']) |
|
|
|
|
| def collate(tokenizer, layout, items): |
| """(model inputs, pooling inputs, padded token count) for a forward over encoded decisions.""" |
| if layout == 'pairs': |
| rows = [r for item in items for r in item] |
| batch = tokenizer.pad(rows, return_tensors='pt') |
| return dict(batch), {}, batch['input_ids'].numel() |
| batch, pool = collate_packed(items, tokenizer.pad_token_id) |
| return batch, pool, batch['input_ids'].numel() |
|
|
|
|
| @torch.inference_mode() |
| def score(model, tokenizer, layout, query, candidates, *, max_length=8192, max_chars=4000, temperature=1.0): |
| """Probabilities (sigmoid of logit / temperature) for one decision's candidates.""" |
| item = encode(tokenizer, layout, query, candidates, max_length, max_chars) |
| batch, pool, _ = collate(tokenizer, layout, [item]) |
| device = next(model.parameters()).device |
| batch = {k: v.to(device) for k, v in batch.items()} |
| pool = {k: v.to(device) if torch.is_tensor(v) else v for k, v in pool.items()} |
| logits = model(**pool, **batch) |
| return torch.sigmoid(logits / temperature).tolist() |
|
|
|
|
| def packing_supported(trunk, tolerance=1e-2): |
| """Whether `trunk` honours the packed layout: a candidate's pooled state must not change |
| when another candidate is added or the order changes.""" |
| try: |
| vocab = trunk.config.get_text_config().vocab_size |
| q, a, b = [5, 6, 7], [8, 9, 10, 11], [12, 13] |
| device = next(trunk.parameters()).device |
| states = [] |
| for cands in ([a, b], [b, a], [a]): |
| batch, pool = collate_packed([pack(q, cands, 1 % vocab, 2 % vocab, 64)], 0) |
| with torch.inference_mode(): |
| hidden = trunk(**{k: v.to(device) for k, v in batch.items()}).last_hidden_state |
| states.append(pooled_spans(hidden, pool['pool_index'].to(device), pool['count'])) |
| ref = states[0][0] |
| scale = ref.abs().max().clamp(min=1e-6) |
| return all(float((s - ref).abs().max() / scale) < tolerance for s in (states[1][1], states[2][0])) |
| except (TypeError, ValueError, RuntimeError, IndexError): |
| return False |
|
|
|
|
| |
| from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel |
|
|
|
|
| class OpenjevScorerConfig(PretrainedConfig): |
| model_type = 'openjev_scorer' |
|
|
| def __init__(self, base_config=None, head='mlp', layout='packed', sequence_tokens=8192, max_chars=4000, |
| temperatures=None, operating_points=None, kind_targets=None, **kwargs): |
| super().__init__(**kwargs) |
| self.base_config, self.head, self.layout = base_config or {}, head, layout |
| |
| self.sequence_tokens, self.max_chars = sequence_tokens, max_chars |
| self.temperatures, self.operating_points = temperatures or {}, operating_points or {} |
| self.kind_targets = kind_targets or {} |
|
|
|
|
| class OpenjevScorer(PreTrainedModel): |
| """A per-candidate scorer exported by `openjev tod export --hf`. `query` is the decision's |
| question, a blank line, then its state (JSON with indent 1 unless it is a string).""" |
| config_class = OpenjevScorerConfig |
| base_model_prefix = 'scorer' |
| _supports_sdpa = True |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| base = dict(config.base_config) |
| trunk_config = AutoConfig.for_model(base.pop('model_type'), **base) |
| trunk = AutoModel.from_config(trunk_config, attn_implementation='sdpa') |
| self.scorer = Scorer(trunk, make_head(config.head, trunk_config.get_text_config().hidden_size)) |
| self.post_init() |
|
|
| def _init_weights(self, module): |
| pass |
|
|
| def forward(self, **inputs): |
| return self.scorer(**inputs) |
|
|
| def score(self, tokenizer, query, candidates, kind=None): |
| """Calibrated probability per candidate (sigmoid; for a Choice kind, a softmax).""" |
| temperature = self.config.temperatures.get(kind, 1.0) if kind else 1.0 |
| if self.config.kind_targets.get(kind) == 'choice': |
| logits = torch.logit(torch.tensor(score(self.scorer, tokenizer, self.config.layout, query, candidates, |
| max_length=self.config.sequence_tokens, |
| max_chars=self.config.max_chars))) |
| return torch.softmax(logits / temperature, 0).tolist() |
| return score(self.scorer, tokenizer, self.config.layout, query, candidates, max_length=self.config.sequence_tokens, |
| max_chars=self.config.max_chars, temperature=temperature) |
|
|
| def keep(self, tokenizer, query, candidates, kind, point=None): |
| """Indices of the candidates to keep at an operating point (default: the model's).""" |
| points = self.config.operating_points |
| name = point or points['default'] |
| value = points['kinds'][kind][name] |
| p = torch.tensor(self.score(tokenizer, query, candidates, kind)) |
| keep = p >= (value if name.startswith('tau') else value * p.max()) |
| keep[int(p.argmax())] |= name.startswith('rho') |
| return [i for i, k in enumerate(keep.tolist()) if k] |
|
|
|
|
| OpenjevScorerConfig.register_for_auto_class() |
| OpenjevScorer.register_for_auto_class('AutoModel') |
|
|
