| """Run the fine-tuned TinyBERT model on new addresses. |
| |
| Standalone script — only needs `transformers`+`torch` (and optionally |
| `gazetteer_lookup.py` + a pincodes CSV for the administrative fields). |
| Unlike the t5/qwen models, this is token classification (BIO tagging), not |
| JSON generation — it always produces a well-formed field dict; there's no |
| "invalid JSON" failure mode to handle. |
| |
| Usage: |
| # Single address |
| python inference_tinybert.py "FLAT NO.32, UTTARA TOWERS, MG ROAD GUWAHATI , Kamrup Unclassified AS 781029" |
| |
| # Batch from stdin (one address per line) |
| cat addresses.txt | python inference_tinybert.py --stdin |
| |
| # Batch from a text file, output JSONL |
| python inference_tinybert.py --file addresses.txt --out results.jsonl |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
|
|
| MAX_LENGTH = 160 |
|
|
| ALL_FIELDS = ( |
| "houseNumber", "houseName", "poi", "street", "subsubLocality", "subLocality", |
| "locality", "village", "subDistrict", "district", "city", "state", "pincode", |
| ) |
|
|
| LABELS = ["O"] + [f"{prefix}-{field}" for field in ALL_FIELDS for prefix in ("B", "I")] |
| ID2LABEL = {i: label for i, label in enumerate(LABELS)} |
|
|
|
|
| def extract_fields(raw_address: str, offsets: list[tuple[int, int]], pred_label_ids: list[int]) -> dict: |
| """Reconstruct the field dict from per-token BIO predictions by slicing the |
| raw address text at each contiguous B-/I- run's char span — never |
| `tokenizer.decode`, which would lose casing and introduce WordPiece-merge |
| artifacts (e.g. "##" continuation glue).""" |
| result = {f: None for f in ALL_FIELDS} |
| current_field = None |
| current_start = None |
| current_end = None |
|
|
| def flush(): |
| if current_field is not None and result[current_field] is None: |
| result[current_field] = raw_address[current_start:current_end] |
|
|
| for (start, end), label_id in zip(offsets, pred_label_ids): |
| if start == end: |
| continue |
| label = ID2LABEL[label_id] |
| if label == "O": |
| flush() |
| current_field = None |
| continue |
| prefix, field = label.split("-", 1) |
| if prefix == "B" or field != current_field: |
| flush() |
| current_field, current_start, current_end = field, start, end |
| else: |
| current_end = end |
| flush() |
| return result |
|
|
|
|
| def load_model(model_id: str): |
| from transformers import AutoModelForTokenClassification, AutoTokenizer |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_id) |
| model = AutoModelForTokenClassification.from_pretrained(model_id) |
| model.eval() |
| return model, tokenizer |
|
|
|
|
| def parse_address(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() |
| result = extract_fields(raw_address, offsets, pred_ids) |
| result["_raw_address"] = raw_address |
| return result |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description="Parse Indian addresses using the fine-tuned TinyBERT model") |
| p.add_argument("address", nargs="?", help="Single address to parse") |
| p.add_argument("--stdin", action="store_true", help="Read addresses from stdin, one per line") |
| p.add_argument("--file", help="Read addresses from a text file") |
| p.add_argument("--out", help="Write JSONL output to file (default: stdout)") |
| p.add_argument("--model", default=".", help="Model dir or HF repo id (default: current directory)") |
| p.add_argument("--pincodes", default=None, |
| help="Path to a pincodes CSV (India Post format) — if given, adds " |
| "districtAdministrative/stateAdministrative/cityAdministrative " |
| "fields via deterministic pincode lookup (does not alter the " |
| "model's own district/state/city)") |
| args = p.parse_args() |
|
|
| import os |
|
|
| model, tokenizer = load_model(args.model) |
|
|
| gazetteer = None |
| if args.pincodes and os.path.exists(args.pincodes): |
| from gazetteer_lookup import load_gazetteer |
|
|
| gazetteer = load_gazetteer(args.pincodes) |
|
|
| if args.address: |
| addresses = [args.address] |
| elif args.stdin: |
| addresses = [line.rstrip("\n") for line in sys.stdin if line.strip()] |
| elif args.file: |
| with open(args.file, encoding="utf-8") as f: |
| addresses = [line.rstrip("\n") for line in f if line.strip()] |
| else: |
| p.print_help() |
| return |
|
|
| out_f = open(args.out, "w", encoding="utf-8") if args.out else None |
| for addr in addresses: |
| result = parse_address(model, tokenizer, addr) |
| if gazetteer is not None: |
| from gazetteer_lookup import add_administrative_fields |
|
|
| raw_addr = result.pop("_raw_address") |
| result = add_administrative_fields(result, gazetteer) |
| result["_raw_address"] = raw_addr |
| line = json.dumps(result, ensure_ascii=False) |
| if out_f: |
| out_f.write(line + "\n") |
| else: |
| print(line) |
|
|
| if out_f: |
| out_f.close() |
| print(f"Wrote {len(addresses):,} results to {args.out}", file=sys.stderr) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
