"""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