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