File size: 2,338 Bytes
5cee4fa | 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 | """Run one offline extract call and compare with the packaged source-oracle fixture."""
import argparse
import json
import os
from pathlib import Path
import sys
sys.dont_write_bytecode = True
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
def main():
assets = Path(__file__).resolve().parents[1] / 'host_assets'
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--seq', type=int, choices=(128, 256, 512), default=None,
help='Validation override; normally the runtime chooses the smallest window')
parser.add_argument('--table', choices=('fp32', 'fp16'), default='fp32')
parser.add_argument('--result', type=Path)
args = parser.parse_args()
sys.path.insert(0, str(assets))
import torch
from runtime import HostRuntime, compare_entities
from runtime.host_runtime import LABELS
torch.set_num_threads(4)
example = json.loads((assets/'example.json').read_text())
with HostRuntime(assets, table=args.table) as host:
if args.seq is None:
entities = host.extract(example['text'], LABELS)
details = host.last_result
else:
details = host.extract_with_details(example['text'], LABELS, seq=args.seq)
entities = details['entities']
comparison = compare_entities(entities, example['entities'])
passed = details['all_finite'] and comparison['span_sets_equal'] and comparison['max_score_diff'] <= 1e-2
capture = details['captured']
result = {'status': 'PASS' if passed else 'FAIL', 'text': example['text'], 'entities': entities,
'comparison': comparison, 'score_atol': 1e-2, 'selected_shape': details['selected_shape'],
'graph_roles': list(details['backends']), 'backends': details['backends'],
'embedding_table': args.table, 'encoded_length': capture['encoded_length'],
'text_word_length': capture['text_word_length'], 'all_finite': details['all_finite'],
'HF_HUB_OFFLINE': os.environ['HF_HUB_OFFLINE']}
if args.result:
args.result.write_text(json.dumps(result, indent=2)+'\n')
print(json.dumps(result, indent=2))
return 0 if passed else 2
if __name__ == '__main__':
raise SystemExit(main())
|