File size: 13,037 Bytes
1d2de8a | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 | """The reference engine: task-agnostic document prefix, float32 arithmetic, chunked prefill.
Contract v2
-----------
prefix = chat template(system=GENERIC, user="Document:\n" + document ...) <- cached once
branch = "\n\n" + task block (instructions, question, options, answer cue) + generation prompt
The document state therefore carries no task instructions: one state serves Boolean and
choice questions. Historical sources (qwen38/, decision_service/, backbone_comparison/) are
not modified.
"""
import copy
import hashlib
import json
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
MODEL = ROOT / 'runtime/qwen38-8bit'
ART = ROOT / 'artifacts/reference'
CONTRACT = 'solomon-document-prefix-v2'
CLASSES = ['yes_only', 'no_only', 'neither', 'both']
LETTERS = 'ABCDEFGH'
CAPTURE_LAYERS = (31, 47, 55) # zero-based decoder layers; final normed state is always captured
CHUNK = 2048
TOKEN_CAP = 40960
SYSTEM = ('You answer questions about the supplied document. Use only the document. '
'Task instructions follow the document; follow them exactly.')
BOOLEAN_TASK = ('Task: classify the evidence for the question using only the document and its explicit rules. '
'A = Yes only. B = No only. C = neither Yes nor No is established. D = both Yes and No are established. '
'A missing fact is not a negative fact. Evidence about another person or subject does not contradict '
'the queried one. Apply explicit time and replacement rules before deciding. '
'Respond with exactly one letter: A, B, C, or D. Do not explain.')
CHOICE_TASK = ('Task: choose the single option that the document best supports. '
'Respond with exactly one letter. Do not explain.')
def sha(path):
return hashlib.sha256(Path(path).read_bytes()).hexdigest()
def dump(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(path.suffix + '.tmp')
tmp.write_text(json.dumps(value, indent=2) + '\n')
tmp.replace(path)
def read_rows(path):
return [json.loads(s) for s in Path(path).read_text().splitlines() if s.strip()]
def softmax(a):
a = np.asarray(a, dtype=np.float64)
p = np.exp(a - a.max())
return p / p.sum()
def boolean_block(question):
return BOOLEAN_TASK + '\nQuestion: ' + question + '\nAnswer (one letter):'
def choice_block(question, options):
if not 2 <= len(options) <= len(LETTERS):
raise ValueError('Choice needs 2-8 options')
lines = [f'{LETTERS[i]}. {text}' for i, text in enumerate(options)]
return CHOICE_TASK + '\nQuestion: ' + question + '\nOptions:\n' + '\n'.join(lines) + '\nAnswer (one letter):'
def block_for(row):
"""Task block and number of answer letters for a benchmark row."""
inp = row['input']
if row.get('task', 'boolean') == 'choice':
return choice_block(inp['question'], inp['options']), len(inp['options'])
return boolean_block(inp['question']), 4
def messages(document, block):
return [{'role': 'system', 'content': SYSTEM},
{'role': 'user', 'content': 'Document:\n' + document + '\n\n' + block}]
class Ledger:
def __init__(self, path, caps=None):
self.path = Path(path)
self.caps = caps or {}
self.counts = json.loads(self.path.read_text()) if self.path.exists() else {
'forwards': 0, 'input_tokens': 0, 'generated_tokens': 0}
def add(self, forwards=0, input_tokens=0, generated_tokens=0):
c = self.counts
c['forwards'] += forwards
c['input_tokens'] += input_tokens
c['generated_tokens'] += generated_tokens
for key, cap in self.caps.items():
if c[key] > cap:
raise RuntimeError(f'Budget cap exceeded: {key} {c[key]} > {cap}')
dump(self.path, c)
class Engine:
def __init__(self, arithmetic='float32', ledger=None, adapter=None):
import mlx.core as mx
from mlx_vlm import load
if arithmetic not in ('float32', 'bfloat16'):
raise ValueError('arithmetic must be float32 or bfloat16')
self.mx = mx
self.arithmetic = arithmetic
self.model, self.processor = load(str(MODEL), trust_remote_code=False, local_files_only=True)
self.model.eval()
self.t = self.processor.tokenizer if hasattr(self.processor, 'tokenizer') else self.processor
self.lm = self.model.language_model
self.adapter = None
if adapter is not None:
from solomon.lora import apply_lora, load_adapter
apply_lora(self.lm)
load_adapter(self.lm, adapter)
self.adapter = str(adapter)
if arithmetic == 'float32':
# Same packed 8-bit weights; floating parameters and activations promoted.
self.model.apply(lambda x: x.astype(mx.float32) if mx.issubdtype(x.dtype, mx.floating) else x)
mx.eval(self.model.parameters())
mx.clear_cache()
self.ledger = ledger
self._letters = {}
# ---- rendering -------------------------------------------------------------------
def render(self, document, block):
text = self.t.apply_chat_template(messages(document, block), tokenize=False,
add_generation_prompt=True, enable_thinking=False)
return text, self.t.encode(text, add_special_tokens=False)
def prefix_ids(self, document):
"""Token ids of the task-agnostic prefix: everything up to the end of the document."""
text, _ = self.render(document, 'X')
marker = '\n\nX'
end = text.rfind(marker)
if end < 0:
raise ValueError('Document boundary missing')
# Leave the boundary token uncached: its tokenisation can depend on what follows.
ids = self.t.encode(text[:end], add_special_tokens=False)[:-1]
if not ids:
raise ValueError('Empty prefix')
return ids
def letter_ids(self, text, ids, n):
out = []
for letter in LETTERS[:n]:
key = letter
if key not in self._letters:
ext = self.t.encode(text + letter, add_special_tokens=False)
if len(ext) != len(ids) + 1 or ext[:-1] != ids:
raise ValueError('Unstable answer-letter continuation')
self._letters[key] = ext[-1]
out.append(self._letters[key])
return out
# ---- execution -------------------------------------------------------------------
def _run(self, ids, cache, offset, capture, chunk=CHUNK):
"""Process ids through the language model. Returns (last logits, {layer: vector})."""
mx = self.mx
if offset + len(ids) > TOKEN_CAP:
raise ValueError(f'{offset + len(ids)} tokens exceeds the {TOKEN_CAP}-token scope cap')
if self.ledger is not None:
self.ledger.add(forwards=1, input_tokens=len(ids))
n = len(ids)
step = n if chunk is None else chunk
if cache is None and step < n:
raise ValueError('Chunked execution needs a cache')
start = 0
logits = feats = None
while start < n:
stop = min(start + step, n)
last = stop == n
piece = mx.array([ids[start:stop]])
pos = mx.broadcast_to(mx.arange(offset + start, offset + stop)[None, None, :], (3, 1, stop - start))
kwargs = dict(cache=cache, position_ids=pos, skip_logits=True, return_hidden=True)
if last and capture:
kwargs['capture_layer_ids'] = list(CAPTURE_LAYERS)
out = self.lm(piece, **kwargs)
hidden = out.hidden_states
if last:
final = hidden[-1][:, -1:, :]
lg = self.lm.lm_head(final)[0, -1].astype(mx.float32)
vecs = {}
if capture:
for layer, h in zip(CAPTURE_LAYERS, hidden[:-1]):
vecs[layer] = h[0, -1].astype(mx.float32)
vecs['final'] = final[0, -1].astype(mx.float32)
mx.eval(lg, *vecs.values())
logits = np.array(lg)
feats = {k: np.array(v) for k, v in vecs.items()}
elif cache is not None:
mx.eval([c.state for c in cache])
del out, hidden
start = stop
if not np.isfinite(logits).all() or any(not np.isfinite(v).all() for v in feats.values()):
raise ValueError('Non-finite model output')
if mx.get_peak_memory() > 100 * 2**30:
raise MemoryError('100 GiB MLX allocation cap exceeded')
return logits, feats
def prefill(self, document, chunk=CHUNK):
"""Encode the task-agnostic document prefix once."""
ids = self.prefix_ids(document)
cache = self.lm.make_cache()
mx = self.mx
if self.ledger is not None:
self.ledger.add(forwards=1, input_tokens=len(ids))
if len(ids) > TOKEN_CAP:
raise ValueError('Prefix exceeds scope cap')
start = 0
while start < len(ids):
stop = min(start + chunk, len(ids))
pos = mx.broadcast_to(mx.arange(start, stop)[None, None, :], (3, 1, stop - start))
self.lm(mx.array([ids[start:stop]]), cache=cache, position_ids=pos, skip_logits=True)
mx.eval([c.state for c in cache])
start = stop
mx.clear_cache()
return {'prefix_ids': ids, 'cache': cache}
def score(self, document, block, n_letters, state=None, capture=True, execution='cached', full_chunk=None):
"""Answer-letter distribution (and features) for one task block."""
text, ids = self.render(document, block)
letters = self.letter_ids(text, ids, n_letters)
fallback = None
if execution == 'cached':
if state is None:
raise ValueError('cached execution needs a prefilled state')
prefix = state['prefix_ids']
if ids[:len(prefix)] != prefix:
execution, fallback = 'full', 'prefix token mismatch'
if execution == 'cached':
branch = copy.deepcopy(state['cache'])
logits, feats = self._run(ids[len(prefix):], branch, len(prefix), capture, chunk=None)
del branch
reused = len(prefix)
elif full_chunk:
# Long prompts: unchunked float32 attention scores would not fit in memory. The whole
# prompt is still processed from scratch, in chunks at different split points.
logits, feats = self._run(ids, self.lm.make_cache(), 0, capture, chunk=full_chunk)
reused = 0
else:
logits, feats = self._run(ids, None, 0, capture, chunk=None)
reused = 0
full = softmax(logits)
self.mx.clear_cache()
return {'letter_logits': logits[letters].astype(np.float64), 'probabilities': softmax(logits[letters]),
'mass': float(full[letters].sum()), 'top_token': int(logits.argmax()),
'top_is_letter': int(logits.argmax()) in letters, 'features': feats,
'execution': execution, 'fallback': fallback, 'prompt_tokens': len(ids),
'reused_prefix_tokens': reused, 'branch_tokens': len(ids) - reused}
def generate(self, document, block, max_tokens=2048, thinking=True):
"""Reasoning-enabled reference. Uses the library generator on the same contract."""
from mlx_vlm import stream_generate
text = self.t.apply_chat_template(messages(document, block), tokenize=False,
add_generation_prompt=True, enable_thinking=thinking)
pieces, last = [], None
for x in stream_generate(self.model, self.processor, prompt=text, max_tokens=max_tokens,
temperature=0., top_p=1., top_k=0, repetition_penalty=1.):
pieces.append(x.text)
last = x
n = getattr(last, 'generation_tokens', 0)
if self.ledger is not None:
self.ledger.add(generated_tokens=n)
return {'text': ''.join(pieces), 'generation_tokens': n,
'finish_reason': getattr(last, 'finish_reason', None)}
def identity(engine):
"""Runtime fingerprint for state handles: no task prompt inside it any more."""
from importlib.metadata import version
payload = {'contract': CONTRACT, 'system_sha256': hashlib.sha256(SYSTEM.encode()).hexdigest(),
'arithmetic': engine.arithmetic, 'adapter': engine.adapter and sha(engine.adapter),
'engine_sha256': sha(Path(__file__)), 'capture_layers': list(CAPTURE_LAYERS),
'model_config_sha256': sha(MODEL / 'config.json'),
'versions': {p: version(p) for p in ('mlx', 'mlx-vlm', 'transformers')}}
payload['fingerprint'] = hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
return payload
|