Spaces:
Running on Zero
Running on Zero
File size: 9,110 Bytes
a0270e2 | 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 | """
Schema definitions, validation, and sub-vocabulary token mapping for parallel constrained decisions.
Supports booleans and categorical enums with cardinality up to 255.
"""
from typing import Dict, Any, List, Tuple, Optional
import numpy as np
class FieldDefinition:
def __init__(self, name: str, field_type: str, description: str, choices: Optional[List[str]] = None):
self.name = name
self.field_type = field_type.lower()
self.description = description
if self.field_type == "boolean":
self.choices = ["true", "false"]
elif self.field_type in ("enum", "choice", "selection"):
if not choices or len(choices) == 0:
raise ValueError(f"Field '{name}' of type enum must have choices defined.")
if len(choices) > 255:
raise ValueError(f"Field '{name}' exceeds maximum cardinality of 255 choices (got {len(choices)}).")
self.choices = choices
else:
raise ValueError(f"Unsupported field type '{field_type}'. Supported types: 'boolean' and 'enum'.")
self.cached_candidate_token_ids: Optional[List[List[int]]] = None
@property
def cardinality(self) -> int:
return len(self.choices)
def compile_candidate_tokens(self, tokenizer):
"""Pre-indexes and caches candidate token IDs so inference runs in microseconds."""
if self.cached_candidate_token_ids is not None:
return self.cached_candidate_token_ids
candidate_tokens_per_choice = []
if self.field_type == "boolean":
true_variants = ['true', ' true', 'True', ' True', 'TRUE', 'yes', ' yes']
true_ids = []
for v in true_variants:
toks = tokenizer.encode(v, add_special_tokens=False)
if toks:
true_ids.append(toks[0])
candidate_tokens_per_choice.append(list(set(true_ids)))
false_variants = ['false', ' false', 'False', ' False', 'FALSE', 'no', ' no']
false_ids = []
for v in false_variants:
toks = tokenizer.encode(v, add_special_tokens=False)
if toks:
false_ids.append(toks[0])
candidate_tokens_per_choice.append(list(set(false_ids)))
else:
for choice in self.choices:
c_clean = str(choice).strip()
variants = [' ' + c_clean, c_clean]
ids = []
for v in variants:
toks = tokenizer.encode(v, add_special_tokens=False)
if toks:
ids.append(toks[0])
candidate_tokens_per_choice.append(list(set(ids)))
self.cached_candidate_token_ids = candidate_tokens_per_choice
return self.cached_candidate_token_ids
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
"type": self.field_type,
"description": self.description,
"choices": self.choices,
"cardinality": self.cardinality,
}
class StructuredSchema:
def __init__(self, schema_dict: Dict[str, Any], tokenizer=None):
self.fields: Dict[str, FieldDefinition] = {}
for field_name, spec in schema_dict.items():
field_type = spec.get("type", "enum")
description = spec.get("description", "")
choices = spec.get("choices", None)
fdef = FieldDefinition(
name=field_name,
field_type=field_type,
description=description,
choices=choices
)
if tokenizer is not None:
fdef.compile_candidate_tokens(tokenizer)
self.fields[field_name] = fdef
def compile_all_tokens(self, tokenizer):
for fdef in self.fields.values():
fdef.compile_candidate_tokens(tokenizer)
def get_field_names(self) -> List[str]:
return list(self.fields.keys())
def __getitem__(self, key: str) -> FieldDefinition:
return self.fields[key]
def __len__(self) -> int:
return len(self.fields)
def to_json_schema_prompt_str(self) -> str:
"""Returns a clean TypeScript/JSON schema representation for naive LLM prompting."""
lines = ["{"]
for name, field in self.fields.items():
if field.field_type == "boolean":
lines.append(f' "{name}": boolean, // {field.description}')
else:
choices_limit = 20 if len(field.choices) > 50 else len(field.choices)
choices_str = " | ".join(f'"{c}"' for c in field.choices[:choices_limit])
if len(field.choices) > choices_limit:
choices_str += f" | ... ({len(field.choices)} total options)"
lines.append(f' "{name}": {choices_str}, // {field.description}')
lines.append("}")
return "\n".join(lines)
def to_parallel_schema_str(self) -> str:
"""Returns a high-density, compact description catalog for minimal prefill token latency."""
lines = []
for name, field in self.fields.items():
desc = field.description.split('\n')[0].strip()
lines.append(f' "{name}": {desc}')
return "\n".join(lines)
to_rlcd_schema_str = to_parallel_schema_str
def compile_parallel_metadata(self, tokenizer):
"""Pre-indexes and caches compact suffixes, token candidate IDs, and common prefixes."""
if hasattr(self, "_parallel_metadata") and self._parallel_metadata is not None:
return self._parallel_metadata
import os
field_items = list(self.fields.items())
suffix_tok_lists = []
suffix_lengths = []
cands_per_field = []
prefixes = []
has_collisions = []
for fname, fdef in field_items:
if fdef.field_type == "boolean":
suffix = f' "{fname}": '
cands = [
tokenizer.encode("true", add_special_tokens=False)[0],
tokenizer.encode("false", add_special_tokens=False)[0]
]
prefix = ""
else:
prefix = os.path.commonprefix(fdef.choices)
suffix = f' "{fname}": "{prefix}'
cands = []
for c in fdef.choices:
rem = c[len(prefix):]
c_toks = tokenizer.encode(rem, add_special_tokens=False)
cands.append(c_toks[0] if c_toks else tokenizer.encode('"', add_special_tokens=False)[0])
toks = tokenizer.encode(suffix, add_special_tokens=False)
suffix_tok_lists.append(toks)
suffix_lengths.append(len(toks))
cands_per_field.append(cands)
prefixes.append(prefix)
has_collisions.append(len(set(cands)) < len(cands))
max_s_len = max(suffix_lengths)
pad_id = tokenizer.pad_token_id or 0
padded = [s + [pad_id] * (max_s_len - len(s)) for s in suffix_tok_lists]
try:
import mlx.core as mx
suffixes_batch = mx.array(padded, dtype=mx.int32)
except Exception:
suffixes_batch = np.array(padded, dtype=np.int32)
self._parallel_metadata = {
"field_items": field_items,
"suffix_lengths": suffix_lengths,
"cands_per_field": cands_per_field,
"prefixes": prefixes,
"has_collisions": has_collisions,
"suffixes_batch": suffixes_batch
}
return self._parallel_metadata
compile_rlcd_metadata = compile_parallel_metadata
def map_candidate_tokens(tokenizer, choices: List[str], is_boolean: bool = False) -> List[List[int]]:
"""Helper fallback when field definition is not pre-compiled."""
f = FieldDefinition("tmp", "boolean" if is_boolean else "enum", "", choices if not is_boolean else None)
return f.compile_candidate_tokens(tokenizer)
def extract_calibrated_probabilities(
next_token_logits: np.ndarray,
candidate_token_ids_list: List[List[int]],
temperature: float = 1.0
) -> Tuple[int, float, List[float]]:
"""
Takes the logits at the decision token position and computes exact
calibrated probabilities across only the constrained candidate choices (K <= 255).
"""
choice_scores = []
for token_ids in candidate_token_ids_list:
if not token_ids:
choice_scores.append(-1e9)
continue
score = max(float(next_token_logits[tid]) for tid in token_ids)
choice_scores.append(score)
scores = np.array(choice_scores, dtype=np.float32) / max(temperature, 1e-4)
shifted = scores - np.max(scores)
exp_scores = np.exp(shifted)
probs = exp_scores / (np.sum(exp_scores) + 1e-12)
winner_idx = int(np.argmax(probs))
winner_prob = float(probs[winner_idx])
return winner_idx, winner_prob, probs.tolist()
|