| """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() |
|
|
