File size: 4,096 Bytes
20d532b | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Evaluate the fine-tuned TinyBERT model on a held-out test set.
Standalone script — only needs `transformers`+`torch`. Expects a JSONL file
where each line is `{"messages": [system, user, assistant]}` (same chat
format used by the training data on GitHub), with the assistant content
being the gold JSON string. See:
https://github.com/innerkorehq/indian-address-parser
Usage:
python evaluate_tinybert.py test.jsonl
python evaluate_tinybert.py test.jsonl --n 100
"""
from __future__ import annotations
import argparse
import json
from collections import defaultdict
from inference_tinybert import ALL_FIELDS, extract_fields, load_model, MAX_LENGTH
def predict(model, tokenizer, raw_address: str) -> dict:
import torch
enc = tokenizer(
raw_address, return_tensors="pt", return_offsets_mapping=True,
truncation=True, max_length=MAX_LENGTH,
)
offsets = enc.pop("offset_mapping")[0].tolist()
with torch.no_grad():
out = model(**enc)
pred_ids = out.logits[0].argmax(-1).tolist()
return extract_fields(raw_address, offsets, pred_ids)
def compute_metrics(pairs: list[tuple[dict, dict]]) -> dict:
n_total = len(pairs)
field_correct = defaultdict(int)
field_present_gold = defaultdict(int)
field_recalled = defaultdict(int)
overall_exact = 0
for gold, pred in pairs:
all_match = True
for field in ALL_FIELDS:
g = (gold.get(field) or "").strip().lower()
p = (pred.get(field) or "").strip().lower()
if g:
field_present_gold[field] += 1
if g == p:
field_recalled[field] += 1
else:
all_match = False
if g == p:
field_correct[field] += 1
if all_match:
overall_exact += 1
accs = {f: (field_correct[f] / n_total if n_total else 0) for f in ALL_FIELDS}
return {
"n_total": n_total,
"overall_exact": overall_exact,
"field_correct": field_correct,
"field_present_gold": field_present_gold,
"field_recalled": field_recalled,
"field_accuracy": accs,
"mean_field_accuracy": sum(accs.values()) / len(accs),
}
def print_report(m: dict):
n_total = m["n_total"]
print(f"\nResults")
print(f"Samples evaluated: {n_total}")
print(f"Overall exact match (all present fields): {m['overall_exact']}/{n_total} ({100*m['overall_exact']/max(1,n_total):.1f}%)\n")
print(f"{'Field':20s} {'Accuracy':>10s} {'Recall':>10s} {'Gold presence':>14s}")
print("-" * 60)
for field in ALL_FIELDS:
acc = m["field_accuracy"][field]
n_gold = m["field_present_gold"][field]
rec = m["field_recalled"][field] / n_gold if n_gold else float("nan")
pres = n_gold / n_total if n_total else 0
rec_str = f"{100*rec:.1f}%" if n_gold else " n/a"
print(f"{field:20s} {100*acc:>9.1f}% {rec_str:>10s} {100*pres:>13.1f}%")
print(f"\n{'Mean field accuracy':20s} {100*m['mean_field_accuracy']:>9.1f}%")
def main():
p = argparse.ArgumentParser()
p.add_argument("test_file", help="JSONL file with {'messages': [system, user, assistant]} per line")
p.add_argument("--model", default=".", help="Model dir or HF repo id (default: current directory)")
p.add_argument("--n", type=int, default=None)
args = p.parse_args()
with open(args.test_file, encoding="utf-8") as f:
samples = [json.loads(line) for line in f]
if args.n:
samples = samples[: args.n]
model, tokenizer = load_model(args.model)
pairs = []
for i, sample in enumerate(samples):
msgs = sample["messages"]
raw_address = msgs[1]["content"].replace("Parse this address:\n", "", 1)
gold = json.loads(msgs[2]["content"])
pred = predict(model, tokenizer, raw_address)
pairs.append((gold, pred))
if (i + 1) % 25 == 0:
print(f" ...{i + 1}/{len(samples)} evaluated")
print_report(compute_metrics(pairs))
if __name__ == "__main__":
main()
|