| """Transformers inference with autoregressive decoding and batched field decisions.""" |
|
|
| import copy |
| import json |
| import os |
| import threading |
| import time |
| from functools import wraps |
| from typing import Any, Dict, Generator |
|
|
| import torch |
| from torch.nn.attention.flex_attention import create_block_mask |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| from core.prompt_builder import build_naive_json_prompt |
| from core.schema import StructuredSchema |
|
|
| MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct") |
| _model = None |
| _tokenizer = None |
| _gpu_lock = threading.Lock() |
| _load_lock = threading.Lock() |
|
|
|
|
| def gpu_locked(fn): |
| @wraps(fn) |
| def wrapper(*args, **kwargs): |
| with _gpu_lock, torch.inference_mode(): |
| return fn(*args, **kwargs) |
| return wrapper |
|
|
|
|
| def gpu_locked_gen(fn): |
| @wraps(fn) |
| def wrapper(*args, **kwargs): |
| |
| |
| iterator = fn(*args, **kwargs) |
| _gpu_lock.acquire() |
| try: |
| while True: |
| with torch.inference_mode(): |
| try: |
| event = next(iterator) |
| except StopIteration: |
| return |
| yield event |
| finally: |
| iterator.close() |
| _gpu_lock.release() |
| return wrapper |
|
|
|
|
| def get_engine(): |
| with _load_lock, torch.inference_mode(): |
| return _load_engine() |
|
|
|
|
| def _load_engine(): |
| global _model, _tokenizer |
| if _model is None: |
| device = torch.device(os.environ.get("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")) |
| dtype = torch.float32 |
| if device.type == "cuda": |
| with torch.cuda.device(device): |
| dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| print(f"Loading {MODEL_ID} on {device} ({dtype})...") |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=dtype).to(device).eval() |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| warmup = tokenizer("Warmup", return_tensors="pt").to(device) |
| model(**warmup, use_cache=True, logits_to_keep=1) |
| _synchronize(model) |
| _model, _tokenizer = model, tokenizer |
| print("Engine ready.") |
| return _model, _tokenizer |
|
|
|
|
| def _synchronize(model): |
| if model.device.type == "cuda": |
| torch.cuda.synchronize(model.device) |
|
|
|
|
| def _naive_result(text, schema, token_count, passes, elapsed_ms): |
| parsed = None |
| error = None |
| try: |
| parsed = json.loads(text.strip()) |
| except ValueError as exc: |
| error = str(exc) |
| valid = error is None |
| missing, invalid = [], [] |
| if isinstance(parsed, dict): |
| for name, field in schema.fields.items(): |
| if name not in parsed: |
| missing.append(name) |
| elif (field.field_type == "boolean" and not isinstance(parsed[name], bool)) or ( |
| field.field_type != "boolean" and parsed[name] not in field.choices |
| ): |
| invalid.append(f"{name}={parsed[name]}") |
| return { |
| "mode": "naive_autoregressive", "elapsed_ms": round(elapsed_ms, 2), |
| "total_tokens": token_count, |
| "tokens_per_second": round(token_count * 1000 / max(elapsed_ms, 1e-6), 1), |
| "sequential_forward_passes": passes, "is_valid_json": valid, |
| "schema_match": isinstance(parsed, dict) and set(parsed) == set(schema.fields) and not invalid, |
| "raw_text": text, "parsed_json": parsed, "parse_error": error, |
| "missing_keys": missing, "invalid_enums": invalid, |
| "has_calibrated_probabilities": False, |
| } |
|
|
|
|
| @gpu_locked_gen |
| def stream_naive_generation(context: str, schema: StructuredSchema, max_tokens: int = 700, |
| temperature: float = 0.2) -> Generator[Dict[str, Any], None, None]: |
| if max_tokens < 1 or temperature < 0: |
| raise ValueError("max_tokens must be positive and temperature nonnegative") |
| model, tokenizer = get_engine() |
| inputs = tokenizer(build_naive_json_prompt(context, schema), add_special_tokens=False, |
| return_tensors="pt").to(model.device) |
| t0 = time.perf_counter() |
| cache = None |
| generated = [] |
| current_text = "{\n " |
| eos = model.generation_config.eos_token_id |
| stop_tokens = set(eos if isinstance(eos, list) else [eos]) | {tokenizer.eos_token_id} |
| passes = 0 |
| emitted = False |
| for _ in range(max_tokens): |
| output = model(**inputs, past_key_values=cache, use_cache=True, logits_to_keep=1) |
| cache = output.past_key_values |
| passes += 1 |
| logits = output.logits[0, -1].float() |
| token = int(torch.argmax(logits) if temperature == 0 else |
| torch.multinomial(torch.softmax(logits / temperature, dim=-1), 1)[0]) |
| if token in stop_tokens: |
| break |
| generated.append(token) |
| text = "{\n " + tokenizer.decode(generated, skip_special_tokens=True, |
| clean_up_tokenization_spaces=False) |
| |
| if not text.endswith("\ufffd"): |
| delta = text[len(current_text):] if emitted else text |
| emitted = True |
| current_text = text |
| yield {"type": "token", "token": delta, "accumulated": text, |
| "token_count": len(generated), |
| "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1)} |
| try: |
| if isinstance(json.loads(text), dict): |
| break |
| except ValueError: |
| pass |
| inputs = {"input_ids": torch.tensor([[token]], device=model.device), |
| "attention_mask": torch.cat([inputs["attention_mask"], |
| torch.ones((1, 1), dtype=torch.long, device=model.device)], dim=1)} |
| _synchronize(model) |
| yield {"type": "done", "result": _naive_result( |
| current_text, schema, len(generated), passes, (time.perf_counter() - t0) * 1000)} |
|
|
|
|
| def run_naive_generation(context: str, schema: StructuredSchema, max_tokens: int = 700, |
| temperature: float = 0.2) -> Dict[str, Any]: |
| for event in stream_naive_generation(context, schema, max_tokens, temperature): |
| if event["type"] == "done": |
| return event["result"] |
|
|
|
|
| def _candidate_scores(model, cache, first_logits, candidates, pad_id): |
| """Score complete allowed continuations, including their closing delimiter.""" |
| first_log_probs = first_logits.float().log_softmax(-1) |
| scores = first_log_probs[[tokens[0] for tokens in candidates]].clone() |
| passes = 0 |
| |
| for start in range(0, len(candidates), 16): |
| sequences = candidates[start:start + 16] |
| width = max(map(len, sequences)) - 1 |
| if width == 0: |
| continue |
| batch = len(sequences) |
| ids = torch.full((batch, width), pad_id, dtype=torch.long, device=model.device) |
| mask = torch.zeros_like(ids) |
| for row, tokens in enumerate(sequences): |
| ids[row, :len(tokens) - 1] = torch.tensor(tokens[:-1], device=model.device) |
| mask[row, :len(tokens) - 1] = 1 |
| branch = copy.deepcopy(cache) |
| branch.batch_repeat_interleave(batch) |
| attention = torch.cat([torch.ones((batch, cache.get_seq_length()), |
| dtype=torch.long, device=model.device), mask], dim=1) |
| output = model(input_ids=ids, attention_mask=attention, |
| past_key_values=branch, use_cache=True) |
| passes += 1 |
| log_probs = output.logits.float().log_softmax(-1) |
| for row, tokens in enumerate(sequences): |
| if len(tokens) > 1: |
| scores[start + row] += log_probs[row, torch.arange(len(tokens) - 1, |
| device=model.device), tokens[1:]].sum() |
| del output, branch, log_probs |
| return scores, passes |
|
|
|
|
| @gpu_locked |
| def run_parallel_generation(context: str, schema: StructuredSchema, |
| temperature: float = 1.0) -> Dict[str, Any]: |
| if temperature < 0: |
| raise ValueError("temperature must be nonnegative") |
| model, tokenizer = get_engine() |
| t0 = time.perf_counter() |
| meta = schema.compile_parallel_metadata(tokenizer) |
| prompt = (f"<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}" |
| f"<|im_end|>\n<|im_start|>user\n{context}<|im_end|>\n" |
| f"<|im_start|>assistant\n{{\n") |
| inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device) |
| t_pre = time.perf_counter() |
| output = model(**inputs, use_cache=True, logits_to_keep=1) |
| cache = output.past_key_values |
| prefix_length = cache.get_seq_length() |
| _synchronize(model) |
| prefill_ms = (time.perf_counter() - t_pre) * 1000 |
| del output |
| cache.batch_repeat_interleave(len(schema)) |
| suffixes = meta["suffixes_batch"].to(model.device) |
| suffix_mask = meta["suffix_attention_mask"].to(model.device) |
| attention = torch.cat([inputs["attention_mask"].repeat(len(schema), 1), suffix_mask], dim=1) |
| t_suffix = time.perf_counter() |
| output = model(input_ids=suffixes, attention_mask=attention, past_key_values=cache, use_cache=True) |
| _synchronize(model) |
| suffix_ms = (time.perf_counter() - t_suffix) * 1000 |
| parsed, telemetry = {}, {} |
| passes = 2 |
| for i, (name, field) in enumerate(meta["field_items"]): |
| length = meta["suffix_lengths"][i] |
| logits = output.logits[i, length - 1] |
| candidates = meta["candidate_sequences"][i] |
| if meta["has_collisions"][i]: |
| branch = copy.deepcopy(cache) |
| branch.batch_select_indices(torch.tensor([i], device=model.device)) |
| |
| padding = branch.get_seq_length() - prefix_length - length |
| if padding: |
| branch.crop(-padding) |
| scores, extra_passes = _candidate_scores(model, branch, logits, candidates, tokenizer.pad_token_id) |
| passes += extra_passes |
| del branch |
| else: |
| scores = logits[[tokens[0] for tokens in candidates]].float() |
| probabilities = torch.softmax(scores / max(temperature, 1e-4), dim=-1).tolist() |
| winner = max(range(len(probabilities)), key=probabilities.__getitem__) |
| choice = field.choices[winner] |
| value = choice == "true" if field.field_type == "boolean" else choice |
| confidence = round(probabilities[winner], 4) |
| parsed[name] = {"value": value, "prob": confidence} |
| top = sorted([{"choice": c, "probability": round(p, 4)} |
| for c, p in zip(field.choices, probabilities)], |
| key=lambda item: item["probability"], reverse=True) |
| telemetry[name] = {"value": value, "type": field.field_type, "confidence": confidence, |
| "cardinality": field.cardinality, "top_choices": top[:5]} |
| _synchronize(model) |
| return {"mode": "parallel_constrained_calibrated", "elapsed_ms": round((time.perf_counter() - t0) * 1000, 2), |
| "prefill_ms": round(prefill_ms, 2), "suffix_eval_ms": round(suffix_ms, 2), |
| "total_tokens_generated": 0, "sequential_forward_passes": passes, |
| "is_valid_json": True, "schema_match": True, "parsed_json": parsed, |
| "field_telemetry": telemetry, "has_calibrated_probabilities": True, "num_fields": len(schema)} |
|
|
|
|
| @torch.compile(fullgraph=True) |
| def _tree_attention_mask(starts, ends, prefix_length): |
| """Compile the mask predicate into block metadata without a dense token mask.""" |
| size = starts.shape[0] |
| limits = torch.tensor([size - 1, prefix_length, prefix_length + size], device=starts.device) |
|
|
| def visible(batch, head, query, key): |
| |
| |
| q = torch.minimum(query.clamp_min(0), limits[0]) |
| k = torch.minimum((key - limits[1]).clamp_min(0), limits[0]) |
| ancestor = (starts[k] <= starts[q]) & (starts[q] < ends[k]) |
| return ((query <= limits[0]) & (key < limits[2]) |
| & ((key < limits[1]) | ancestor)) |
|
|
| return create_block_mask(visible, B=None, H=None, Q_LEN=size, |
| KV_LEN=prefix_length + size, device=starts.device) |
|
|
|
|
| @gpu_locked |
| def run_parallel_generation_tree(context: str, schema: StructuredSchema, |
| temperature: float = 1.0) -> Dict[str, Any]: |
| """Score all complete candidates in one tree pass over a single prefix cache. |
| |
| Uses a FlexAttention BlockMask with DFS intervals for ancestor visibility. |
| Scores are mean token log-probabilities over the full candidate, including |
| shared value prefixes and closing delimiters. The model's original |
| attention backend is restored afterward. |
| """ |
| if temperature < 0: |
| raise ValueError("temperature must be nonnegative") |
| model, tokenizer = get_engine() |
| if "sliding_attention" in getattr(model.config, "layer_types", []): |
| raise ValueError("Tree generation requires full attention, not a sliding-window cache") |
| t0 = time.perf_counter() |
| meta = schema.compile_tree_metadata(tokenizer) |
| prompt = (f"<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}" |
| f"<|im_end|>\n<|im_start|>user\n{context}<|im_end|>\n" |
| f"<|im_start|>assistant\n{{\n") |
| inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device) |
| t_pre = time.perf_counter() |
| output = model(**inputs, use_cache=True, logits_to_keep=1) |
| cache = output.past_key_values |
| prefix_length = cache.get_seq_length() |
| _synchronize(model) |
| prefill_ms = (time.perf_counter() - t_pre) * 1000 |
| del output |
|
|
| t_suffix = time.perf_counter() |
| ids = meta["input_ids"].to(model.device) |
| positions = meta["depths"].to(model.device) + prefix_length |
| size = ids.shape[1] |
| attention = _tree_attention_mask(meta["dfs_starts"].to(model.device), |
| meta["dfs_ends"].to(model.device), prefix_length) |
| original_backend = model.config._attn_implementation |
| try: |
| model.set_attn_implementation("flex_attention") |
| |
| |
| output = model(input_ids=ids, attention_mask=attention, position_ids=positions, |
| past_key_values=cache, use_cache=True, |
| logits_to_keep=meta["logit_indices"].to(model.device), |
| kernel_options={"FORCE_USE_FLEX_ATTENTION": True}) |
| finally: |
| model.set_attn_implementation(original_backend) |
| _synchronize(model) |
| suffix_ms = (time.perf_counter() - t_suffix) * 1000 |
| |
| |
| logits = output.logits[0].float() |
| rows = meta["score_rows"].to(model.device) |
| tokens = meta["score_tokens"].to(model.device) |
| edge_scores = logits[rows, tokens] - logits.logsumexp(-1)[rows] |
| scores = logits.new_zeros(sum(meta["candidate_counts"])) |
| scores.scatter_add_(0, meta["score_candidates"].to(model.device), edge_scores) |
| scores /= meta["candidate_lengths"].to(model.device) |
| del output, cache, logits, attention |
|
|
| parsed, telemetry = {}, {} |
| for (name, field), field_scores in zip(meta["field_items"], scores.split(meta["candidate_counts"])): |
| probabilities = torch.softmax(field_scores / max(temperature, 1e-4), dim=-1).tolist() |
| winner = max(range(len(probabilities)), key=probabilities.__getitem__) |
| choice = field.choices[winner] |
| value = choice == "true" if field.field_type == "boolean" else choice |
| confidence = round(probabilities[winner], 4) |
| parsed[name] = {"value": value, "prob": confidence} |
| top = sorted([{"choice": c, "probability": round(p, 4)} |
| for c, p in zip(field.choices, probabilities)], |
| key=lambda item: item["probability"], reverse=True) |
| telemetry[name] = {"value": value, "type": field.field_type, "confidence": confidence, |
| "cardinality": field.cardinality, "top_choices": top[:5]} |
| _synchronize(model) |
| return {"mode": "parallel_constrained_tree", "elapsed_ms": round((time.perf_counter() - t0) * 1000, 2), |
| "prefill_ms": round(prefill_ms, 2), "suffix_eval_ms": round(suffix_ms, 2), |
| "total_tokens_generated": 0, "sequential_forward_passes": 2, |
| "is_valid_json": True, "schema_match": True, "parsed_json": parsed, |
| "field_telemetry": telemetry, "has_calibrated_probabilities": True, "num_fields": len(schema), |
| "tree_nodes": size, "prefix_tokens": prefix_length, |
| "scoring": "mean_token_log_probability"} |
|
|
|
|
| run_rlcd_generation = run_parallel_generation |
|
|
