| """Save unfiltered predictions from the deployed ONNX model for evaluation.""" |
| import argparse |
| from collections import Counter |
| import json |
| import sys |
| import time |
| from pathlib import Path |
| sys.path.insert(0,str(Path(__file__).resolve().parents[1])) |
| import numpy as np |
| import onnxruntime as ort |
| from transformers import AutoTokenizer |
| from pii.core import decode_window, result, encode_windows |
|
|
| def main(): |
| p=argparse.ArgumentParser() |
| p.add_argument("--bundle",type=Path,default=Path("artifacts/bundle")) |
| p.add_argument("--input",type=Path,required=True) |
| p.add_argument("--output",type=Path,required=True) |
| p.add_argument("--batch-size",type=int,default=16,help="Measured training recipe batch; dynamically padded ONNX batches") |
| args=p.parse_args() |
| config=json.loads((args.bundle/"config.json").read_text()) |
| tokenizer=AutoTokenizer.from_pretrained(args.bundle,local_files_only=True) |
| options=ort.SessionOptions();options.intra_op_num_threads=1 |
| session=ort.InferenceSession(str(args.bundle/"model.onnx"),sess_options=options,providers=["CPUExecutionProvider"]) |
| args.output.parent.mkdir(parents=True,exist_ok=True) |
| started=time.perf_counter();documents=windows=0 |
| def batch_predict(rows,handle): |
| nonlocal documents,windows |
| encoded={"input_ids":[],"attention_mask":[],"offset_mapping":[]} |
| mapping=[] |
| for i,row in enumerate(rows): |
| windows_for_row=encode_windows(tokenizer,row["text"],config["max_length"],config["stride"]) |
| for key in encoded:encoded[key].extend(windows_for_row[key]) |
| mapping.extend([i]*len(windows_for_row["input_ids"])) |
| counts=Counter(mapping);seen=Counter();findings=[[] for _ in rows] |
| for begin in range(0,len(mapping),args.batch_size): |
| end=min(len(mapping),begin+args.batch_size) |
| length=max(len(x) for x in encoded["input_ids"][begin:end]) |
| ids=np.full((end-begin,length),tokenizer.pad_token_id,dtype=np.int64) |
| masks=np.zeros_like(ids) |
| for j,index in enumerate(range(begin,end)): |
| n=len(encoded["input_ids"][index]);ids[j,:n]=encoded["input_ids"][index];masks[j,:n]=encoded["attention_mask"][index] |
| logits=session.run(None,{"input_ids":ids,"attention_mask":masks})[0] |
| for j,index in enumerate(range(begin,end)): |
| sample=mapping[index];offsets=encoded["offset_mapping"][index] |
| findings[sample].extend(decode_window(logits[j,:len(offsets)],offsets,rows[sample]["text"],config["max_length"],has_previous=seen[sample]>0,has_next=seen[sample]+1<counts[sample])) |
| seen[sample]+=1 |
| for row,found in zip(rows,findings): |
| prediction={"id":row["id"],"text":row["text"],**result(row["text"],found)} |
| handle.write(json.dumps(prediction,ensure_ascii=False)+"\n") |
| documents+=len(rows);windows+=len(mapping) |
| args.output.with_suffix(".progress.json").write_text(json.dumps({"documents":documents,"windows":windows,"seconds":time.perf_counter()-started})+"\n") |
| temporary=args.output.with_suffix(".partial.jsonl") |
| with args.input.open() as source,temporary.open("w") as target: |
| rows=[] |
| for line in source: |
| if not line.strip():continue |
| rows.append(json.loads(line)) |
| if len(rows)==args.batch_size:batch_predict(rows,target);rows=[] |
| if rows:batch_predict(rows,target) |
| temporary.replace(args.output) |
| evidence={"documents":documents,"windows":windows,"seconds":time.perf_counter()-started,"backend":"ONNX Runtime CPU","batch_size":args.batch_size,"raw_unfiltered":True} |
| args.output.with_suffix(".evidence.json").write_text(json.dumps(evidence,indent=2)+"\n") |
| print(json.dumps(evidence),flush=True) |
|
|
| if __name__=="__main__":main() |
|
|
