| """Drop-in PyTorch SDK adapter for independent-field tree attention."""
|
| import time
|
|
|
| import torch
|
|
|
| from core.engine_torch import get_torch_engine, gpu_decorator
|
| from tree_decode import FieldPlan, decode_fields, prefill
|
|
|
|
|
| @gpu_decorator
|
| @torch.inference_mode()
|
| def run_parallel_generation_tree(context, schema, temperature=1.0):
|
| model, tokenizer, device = get_torch_engine()
|
| if model.config._attn_implementation not in ('sdpa', 'eager'):
|
| raise ValueError('Tree attention requires SDPA or eager attention with a 4D mask')
|
|
|
| def sync():
|
| if str(device).startswith('mps'):
|
| torch.mps.synchronize()
|
|
|
| sync()
|
| start = time.perf_counter()
|
| meta = schema.compile_parallel_metadata(tokenizer)
|
| suffixes = [list(map(int, row[:n])) for row, n in zip(meta['suffixes_batch'], meta['suffix_lengths'])]
|
| prompt = (f'<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}<|im_end|>\n'
|
| f'<|im_start|>user\n{context}<|im_end|>\n<|im_start|>assistant\n{{\n')
|
| ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
|
| plan = FieldPlan(suffixes, ids.shape[1], device, model.dtype, tokenizer.pad_token_id or 0)
|
| sync()
|
| pre_start = time.perf_counter()
|
| cache = prefill(model, ids)
|
| sync()
|
| pre_ms = (time.perf_counter() - pre_start) * 1000
|
| dec_start = time.perf_counter()
|
| logits = decode_fields(model, cache, plan, 'tree')[0]
|
| sync()
|
| dec_ms = (time.perf_counter() - dec_start) * 1000
|
| parsed, telemetry = {}, {}
|
| for i, (name, field) in enumerate(meta['field_items']):
|
| probs = (logits[i, meta['cands_per_field'][i]].float() / max(temperature, 1e-4)).softmax(-1).tolist()
|
| winner = max(range(len(probs)), key=probs.__getitem__)
|
| value = winner == 0 if field.field_type == 'boolean' else field.choices[winner]
|
| parsed[name] = {'value': value, 'prob': round(probs[winner], 4)}
|
| choices = [{'choice': c, 'probability': round(p, 4)} for c, p in zip(field.choices, probs)]
|
| telemetry[name] = dict(value=value, type=field.field_type, confidence=round(probs[winner], 4),
|
| cardinality=field.cardinality, top_choices=sorted(choices, key=lambda c: c['probability'], reverse=True)[:5])
|
| return dict(mode='parallel_constrained_tree', elapsed_ms=round((time.perf_counter()-start)*1000, 2),
|
| prefill_ms=round(pre_ms, 2), suffix_eval_ms=round(dec_ms, 2), total_tokens_generated=0,
|
| sequential_forward_passes=1, is_valid_json=True, schema_match=True,
|
| parsed_json=parsed, field_telemetry=telemetry, has_calibrated_probabilities=False,
|
| num_fields=len(schema), device=str(device),
|
| candidate_collision_fields=[name for (name, _), collision in zip(meta['field_items'], meta['has_collisions']) if collision])
|
|
|
