| """Run the upstream field classifier with shared-prefix tree attention."""
|
| import argparse
|
| import json
|
| import sys
|
| from pathlib import Path
|
|
|
| import torch
|
| from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
| from tree_decode import FieldPlan, decode_fields, prefill
|
|
|
| ROOT = Path(__file__).resolve().parent
|
| SOURCE = ROOT / 'upstream' if (ROOT / 'upstream').exists() else ROOT
|
| sys.path.insert(0, str(SOURCE))
|
| from core.schema import StructuredSchema
|
|
|
|
|
| @torch.inference_mode()
|
| def classify(model, tokenizer, context, schema, mode='tree'):
|
| """Match upstream's one-token candidate scoring, including its limitations.
|
|
|
| Returned probabilities are normalized candidate scores, not empirically
|
| calibrated confidence. Duplicate candidate tokens are explicitly reported.
|
| """
|
| device, dtype = model.device, model.dtype
|
| meta = schema.compile_parallel_metadata(tokenizer)
|
| suffixes = [list(map(int, row[:length])) for row, length 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, dtype, tokenizer.pad_token_id or 0)
|
| cache = prefill(model, ids)
|
| logits = decode_fields(model, cache, plan, mode)[0]
|
| fields = {}
|
| for i, (name, field) in enumerate(meta['field_items']):
|
| probs = logits[i, meta['cands_per_field'][i]].float().softmax(-1)
|
| winner = probs.argmax().item()
|
| fields[name] = {'value': winner == 0 if field.field_type == 'boolean' else field.choices[winner],
|
| 'candidate_probability': probs[winner].item()}
|
| return {'mode': mode, 'fields': fields,
|
| 'candidate_collision_fields': [name for (name, _), collision in zip(meta['field_items'], meta['has_collisions']) if collision]}
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument('--preset', default='fintech_fraud')
|
| ap.add_argument('--mode', choices=['tree', 'batch'], default='tree')
|
| ap.add_argument('--device', choices=['auto', 'mps', 'cpu'], default='auto')
|
| args = ap.parse_args()
|
| torch.set_num_threads(4)
|
| name = 'Qwen/Qwen2.5-1.5B-Instruct'
|
| tokenizer = AutoTokenizer.from_pretrained(name)
|
| device = args.device
|
| if device == 'auto':
|
| device = 'mps' if torch.backends.mps.is_available() else 'cpu'
|
| model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16 if device == 'mps' else torch.float32,
|
| attn_implementation='sdpa').to(device).eval()
|
| preset = json.loads((SOURCE / 'presets' / f'{args.preset}.json').read_text(encoding='utf-8'))
|
| result = classify(model, tokenizer, preset['context'], StructuredSchema(preset['schema']), args.mode)
|
| print(json.dumps(result, indent=2))
|
|
|
|
|
| if __name__ == '__main__':
|
| main()
|
|
|
