File size: 7,820 Bytes
53e24ca e099c73 53e24ca f092803 53e24ca f092803 53e24ca f092803 53e24ca e099c73 53e24ca e099c73 53e24ca e099c73 53e24ca e099c73 53e24ca e099c73 53e24ca f092803 53e24ca f092803 53e24ca e099c73 53e24ca | 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 | """Score native JSON values against the same multimodal prompt as normal Gemma."""
import math
import time
from collections import OrderedDict
from .cached_backend import CachedMLXBackend, PreparedState
from .comparison import prepare_generation
from .core import ScoringRequest, TokenScores
from .json_scoring import candidate_fields, compile_field
class JSONMLXBackend(CachedMLXBackend):
probability_source = "restricted_json_value_likelihoods"
def _sequence_scores(self, prefix_cache, prefix_tokens, fields, on_scores=None):
"""Teacher-force every complete candidate in bounded GPU batches.
Score all tokens, including string terminators. Shared first tokens and
prefix labels never fall back to an arbitrary choice or synthetic confidence.
"""
mx = self.mx
jobs = [
(field_index, choice, field.prefix, candidate)
for field_index, field in fields
for choice, candidate in enumerate(field.candidates)
]
values = {index: [None] * len(field.candidates) for index, field in fields}
batches = []
output = {}
for start in range(0, len(jobs), self.branch_batch_size):
batch = jobs[start : start + self.branch_batch_size]
suffixes = [prefix + candidate[:-1] for _, _, prefix, candidate in batch]
width = max(map(len, suffixes))
tokens = mx.array(
[
list(suffix) + [self.tokenizer.pad_token_id] * (width - len(suffix))
for suffix in suffixes
]
)
tail = width - min(len(prefix) for _, _, prefix, _ in batch) + 1
hidden = self.model.language_model.model(
inputs=tokens, cache=self.fork_cache(prefix_cache, len(batch)), logits_to_keep=tail
)
totals = []
for row, (_, _, prefix, candidate) in enumerate(batch):
position = len(prefix) - 1 - (width - hidden.shape[1])
states = hidden[row : row + 1, position : position + len(candidate), :]
logits = self.model.language_model.logits_from_hidden(states)[0].astype(mx.float32)
log_probs = logits[mx.arange(len(candidate)), mx.array(candidate)] - mx.logsumexp(
logits, axis=-1
)
totals.append(mx.sum(log_probs))
mx.eval(*totals)
for (field_index, choice, _, _), total in zip(batch, totals, strict=True):
values[field_index][choice] = float(total.item())
batches.append(len(batch))
completed = []
for index, field in fields:
if index in output or any(value is None for value in values[index]):
continue
scores = tuple(values[index])
mass = min(1.0, sum(math.exp(value) for value in scores))
length = prefix_tokens + len(field.prefix) + max(map(len, field.candidates)) - 1
output[index] = TokenScores(scores, mass, length)
completed.append((index, output[index]))
if on_scores is not None and completed:
on_scores(completed)
return output, batches
def score_questions(self, state, questions, on_scores=None, on_progress=None):
started = time.perf_counter()
if on_progress:
on_progress("preparing", None)
prompt, inputs = prepare_generation(self, state, questions)
candidates = tuple(candidate_fields(questions))
key = (prompt, candidates)
if not hasattr(self, "_field_cache"):
self._field_cache = OrderedDict()
fields = self._field_cache.get(key)
schema_cache_hit = fields is not None
if fields is None:
base_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
fields = [
compile_field(self.tokenizer, field, prompt, base_ids=base_ids)
for field in candidates
]
self._field_cache[key] = fields
if len(self._field_cache) > 8:
self._field_cache.popitem(last=False)
else:
self._field_cache.move_to_end(key)
prefix_tokens = int(inputs["input_ids"].shape[1])
if any(
prefix_tokens + len(field.prefix) + max(map(len, field.candidates))
> self.max_input_tokens
for field in fields
):
raise ValueError(
"State, questions, and JSON candidate exceed the token limit; no truncation applied"
)
prepared = PreparedState(inputs, [], prefix_tokens)
processed = time.perf_counter()
if on_progress:
on_progress("prefill", prefix_tokens)
prefix_cache = self.prefill(prepared)
prefilled = time.perf_counter()
if on_progress:
on_progress("scoring", prefix_tokens)
self.last_stats = {}
scores = [None] * len(fields)
single = [
(index, field)
for index, field in enumerate(fields)
if all(len(candidate) == 1 for candidate in field.candidates)
]
batches = []
if single:
requests = []
for _, field in single:
symbols = tuple(f"json-token:{candidate[0]}" for candidate in field.candidates)
self._ids.update(
{
symbol: candidate[0]
for symbol, candidate in zip(symbols, field.candidates, strict=True)
}
)
requests.append(ScoringRequest("", symbols))
simple = PreparedState(
inputs, [list(field.prefix) for _, field in single], prefix_tokens
)
def completed_batch(start, results):
on_scores(
[(single[start + offset][0], score) for offset, score in enumerate(results)]
)
results = self.branches(
simple, prefix_cache, requests, on_batch=completed_batch if on_scores else None
)
batches.extend(self.last_stats["branch_batch_sizes"])
for (index, _), result in zip(single, results, strict=True):
scores[index] = result
multiple = [
(index, field)
for index, field in enumerate(fields)
if any(len(candidate) != 1 for candidate in field.candidates)
]
candidate_batches = []
if multiple:
results, candidate_batches = self._sequence_scores(
prefix_cache, prefix_tokens, multiple, on_scores=on_scores
)
for index, result in results.items():
scores[index] = result
self.last_stats = {
"execution": "shared_json_prefix_gpu_batched_fields",
"prefix_prefills": 1,
"prefix_tokens": prefix_tokens,
"schema_cache_hit": schema_cache_hit,
"primitive_fields": len(fields),
"question_suffix_tokens": [len(field.prefix) for field in fields],
"candidate_token_lengths": [
[len(candidate) for candidate in field.candidates] for field in fields
],
"preprocess_seconds": processed - started,
"prefill_seconds": prefilled - processed,
"branch_seconds": time.perf_counter() - prefilled,
"branch_batch_sizes": batches,
"candidate_batch_sizes": candidate_batches,
"compute_dtype": "float32",
"kv_storage": "replicated_per_batch_row",
"conditioning": "shared complete question schema; no previous field answers",
}
return scores
|