pennaburry/parallel-constrained-decoding

1"""2PyTorch / CUDA / CPU Inference Engine for Parallel Constrained Decoding.3Optimized for Linux containers, Hugging Face Spaces (ZeroGPU & CUDA), and cloud environments.4"""5 6import os7import time8import json9import copy10import re11import threading12from typing import Dict, Any, Generator, Optional, List, Tuple13 14import torch15import torch.nn.functional as F16from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer17from transformers.cache_utils import DynamicCache18 19from core.schema import StructuredSchema20from core.prompt_builder import build_naive_json_prompt21 22MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")23 24_torch_model = None25_torch_tokenizer = None26_torch_device = None27_gpu_lock = threading.Lock()28 29# Support Hugging Face Spaces ZeroGPU if available30try:31    import spaces32    gpu_decorator = spaces.GPU(duration=60)33except Exception:34    def gpu_decorator(fn=None, **kwargs):35        if fn is not None:36            return fn37        return lambda f: f38 39 40def get_torch_engine():41    global _torch_model, _torch_tokenizer, _torch_device42    if _torch_model is None or _torch_tokenizer is None:43        _torch_device = "cuda" if torch.cuda.is_available() else "cpu"44        45        if _torch_device == "cuda":46            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float1647        else:48            dtype = torch.float3249 50        print(f"Loading {MODEL_ID} on {_torch_device} ({dtype})...")51        t0 = time.perf_counter()52        _torch_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)53        54        load_kwargs = {55            "torch_dtype": dtype,56            "low_cpu_mem_usage": True57        }58        if _torch_device == "cuda":59            load_kwargs["device_map"] = "auto"60 61        _torch_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **load_kwargs)62        if _torch_device == "cpu":63            _torch_model = _torch_model.to("cpu")64        _torch_model.eval()65        print(f"Engine loaded on {_torch_device} in {time.perf_counter() - t0:.2f}s.")66        67    return _torch_model, _torch_tokenizer, _torch_device68 69 70@gpu_decorator71def run_parallel_generation_torch(72    context: str,73    schema: StructuredSchema,74    temperature: float = 1.075) -> Dict[str, Any]:76    """77    Parallel Constrained Decision Engine running on PyTorch (CUDA / CPU).78    Evaluates all schema fields concurrently against a broadcast prefix KV-cache.79    """80    model, tokenizer, device = get_torch_engine()81    t0 = time.perf_counter()82 83    # 1. Compile schema metadata84    meta = schema.compile_parallel_metadata(tokenizer)85    field_items = meta["field_items"]86    suffix_lengths = meta["suffix_lengths"]87    cands_per_field = meta["cands_per_field"]88    prefixes = meta["prefixes"]89    has_collisions = meta["has_collisions"]90    suffixes_batch = meta["suffixes_batch"]91    M = len(field_items)92 93    # 2. High-density semantic catalog prefill94    schema_str = schema.to_parallel_schema_str()95    base_prompt = (96        f"<|im_start|>system\n"97        f"Classify JSON attributes:\n{schema_str}<|im_end|>\n"98        f"<|im_start|>user\n"99        f"{context}<|im_end|>\n"100        f"<|im_start|>assistant\n{{\n"101    )102    base_toks = tokenizer.encode(base_prompt, return_tensors="pt").to(device)103 104    t_pre0 = time.perf_counter()105    with torch.no_grad():106        base_out = model(base_toks, use_cache=True)107        base_cache = base_out.past_key_values108    t_prefill = (time.perf_counter() - t_pre0) * 1000109 110    # 3. Parallel Suffix Evaluation111    t_suf0 = time.perf_counter()112    pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id or 0113 114    suffix_arr = torch.tensor(suffixes_batch, dtype=torch.long, device=device)115    suffix_mask = (suffix_arr != pad_id).long()116 117    # Broadcast KV cache to batch size M118    with torch.no_grad():119        batched_cache = copy.deepcopy(base_cache)120        if hasattr(batched_cache, "batch_repeat_interleave"):121            batched_cache.batch_repeat_interleave(M)122        elif isinstance(batched_cache, tuple):123            batched_cache = tuple(124                tuple(t.repeat(M, 1, 1, 1) for t in layer)125                for layer in batched_cache126            )127 128        prefix_len = base_toks.shape[1]129        prefix_mask = torch.ones((M, prefix_len), dtype=torch.long, device=device)130        full_mask = torch.cat([prefix_mask, suffix_mask], dim=1)131 132        out = model(suffix_arr, past_key_values=batched_cache, attention_mask=full_mask)133        suffix_out = out.logits134 135    t_suffix_eval = (time.perf_counter() - t_suf0) * 1000136 137    # 4. Slicing, Disambiguation & Softmax138    parsed_json = {}139    field_telemetry = {}140 141    for i, (fname, fdef) in enumerate(field_items):142        decision_idx = suffix_lengths[i] - 1143        field_logits = suffix_out[i, decision_idx, :]144        cand_tokens = cands_per_field[i]145 146        scores = [float(field_logits[tid].item()) for tid in cand_tokens]147        scores_t = torch.tensor(scores, dtype=torch.float32) / max(temperature, 1e-4)148        probs = F.softmax(scores_t, dim=-1).tolist()149        w_idx = int(torch.argmax(scores_t).item())150        w_prob = float(probs[w_idx])151        all_probs = probs152 153        if fdef.field_type == "boolean":154            val = (w_idx == 0)155        else:156            val = fdef.choices[w_idx]157 158        parsed_json[fname] = {159            "value": val,160            "prob": round(w_prob, 4)161        }162 163        choices_list = ["true", "false"] if fdef.field_type == "boolean" else fdef.choices164        scored_choices = []165        for c, p in zip(choices_list, all_probs):166            scored_choices.append({"choice": c, "probability": round(p, 4)})167        scored_choices.sort(key=lambda x: x["probability"], reverse=True)168 169        field_telemetry[fname] = {170            "value": val,171            "type": fdef.field_type,172            "confidence": round(w_prob, 4),173            "cardinality": fdef.cardinality,174            "top_choices": scored_choices[:5]175        }176 177    total_elapsed_ms = (time.perf_counter() - t0) * 1000178 179    return {180        "mode": "parallel_constrained_calibrated",181        "elapsed_ms": round(total_elapsed_ms, 2),182        "prefill_ms": round(t_prefill, 2),183        "suffix_eval_ms": round(t_suffix_eval, 2),184        "total_tokens_generated": 0,185        "sequential_forward_passes": 1,186        "is_valid_json": True,187        "schema_match": True,188        "parsed_json": parsed_json,189        "field_telemetry": field_telemetry,190        "has_calibrated_probabilities": True,191        "num_fields": len(schema),192        "device": device193    }194 195 196@gpu_decorator197def run_naive_generation_torch(198    context: str,199    schema: StructuredSchema,200    temperature: float = 0.2,201    max_new_tokens: int = 512202) -> Dict[str, Any]:203    """204    Standard autoregressive baseline using PyTorch.205    """206    model, tokenizer, device = get_torch_engine()207    t0 = time.perf_counter()208 209    prompt = build_naive_json_prompt(context, schema)210    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)211    prompt_tokens = input_ids.shape[1]212 213    with torch.no_grad():214        output_ids = model.generate(215            input_ids,216            max_new_tokens=max_new_tokens,217            do_sample=(temperature > 0.0),218            temperature=max(temperature, 1e-4),219            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id220        )221 222    elapsed_ms = (time.perf_counter() - t0) * 1000223    gen_tokens = output_ids.shape[1] - prompt_tokens224    tok_per_sec = (gen_tokens / (elapsed_ms / 1000.0)) if elapsed_ms > 0 else 0.0225 226    raw_text = tokenizer.decode(output_ids[0][prompt_tokens:], skip_special_tokens=True)227 228    # Parse JSON229    parsed_json = None230    is_valid = False231    try:232        first_brace = raw_text.find("{")233        last_brace = raw_text.rfind("}")234        if first_brace != -1 and last_brace != -1:235            cleaned = raw_text[first_brace:last_brace + 1]236            parsed_json = json.loads(cleaned)237            is_valid = True238    except Exception:239        pass240 241    schema_match = False242    if is_valid and isinstance(parsed_json, dict):243        expected_keys = set(schema.get_field_names())244        schema_match = (set(parsed_json.keys()) == expected_keys)245 246    return {247        "mode": "autoregressive_naive",248        "elapsed_ms": round(elapsed_ms, 2),249        "total_tokens": gen_tokens,250        "tokens_per_second": round(tok_per_sec, 1),251        "sequential_forward_passes": gen_tokens,252        "is_valid_json": is_valid,253        "schema_match": schema_match,254        "raw_text": raw_text,255        "parsed_json": parsed_json,256        "device": device257    }258 259 260def stream_naive_generation_torch(261    context: str,262    schema: StructuredSchema,263    temperature: float = 0.2,264    max_new_tokens: int = 512265) -> Generator[Dict[str, Any], None, None]:266    """267    Generator streaming individual tokens for side-by-side comparison visualizer.268    """269    model, tokenizer, device = get_torch_engine()270    t0 = time.perf_counter()271 272    prompt = build_naive_json_prompt(context, schema)273    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)274    prompt_tokens = input_ids.shape[1]275 276    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)277    278    gen_kwargs = {279        "input_ids": input_ids,280        "max_new_tokens": max_new_tokens,281        "do_sample": (temperature > 0.0),282        "temperature": max(temperature, 1e-4),283        "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,284        "streamer": streamer285    }286 287    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)288    thread.start()289 290    full_text = ""291    tok_count = 0292 293    for token_str in streamer:294        tok_count += 1295        full_text += token_str296        yield {297            "type": "token",298            "token": token_str,299            "token_count": tok_count300        }301 302    thread.join()303    elapsed_ms = (time.perf_counter() - t0) * 1000304 305    parsed_json = None306    is_valid = False307    try:308        first_brace = full_text.find("{")309        last_brace = full_text.rfind("}")310        if first_brace != -1 and last_brace != -1:311            cleaned = full_text[first_brace:last_brace + 1]312            parsed_json = json.loads(cleaned)313            is_valid = True314    except Exception:315        pass316 317    schema_match = False318    if is_valid and isinstance(parsed_json, dict):319        expected_keys = set(schema.get_field_names())320        schema_match = (set(parsed_json.keys()) == expected_keys)321 322    result = {323        "mode": "autoregressive_naive",324        "elapsed_ms": round(elapsed_ms, 2),325        "total_tokens": tok_count,326        "tokens_per_second": round((tok_count / (elapsed_ms / 1000.0)) if elapsed_ms > 0 else 0.0, 1),327        "sequential_forward_passes": tok_count,328        "is_valid_json": is_valid,329        "schema_match": schema_match,330        "raw_text": full_text,331        "parsed_json": parsed_json,332        "device": device333    }334 335    yield {336        "type": "done",337        "result": result338    }339