File size: 3,078 Bytes
f36843f 728caeb f36843f 728caeb f36843f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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()
|