Spaces:
Sleeping
Sleeping
| """Backend comparison: torch fp32 vs ONNX fp32 vs ONNX int8. | |
| Answers the two deployment questions directly, with measurements: | |
| Q1 Does ONNX Runtime beat PyTorch on CPU for this model? | |
| Q2 Does int8 quantization change the answers? | |
| Accuracy is compared against ground truth, not just "do the two agree", because | |
| two backends can agree and both be wrong. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import statistics | |
| import sys | |
| import time | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parent.parent | |
| sys.path.insert(0, str(ROOT / "src")) | |
| os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") | |
| import numpy as np | |
| import torch | |
| from docxextract import parsing | |
| from docxextract.config import Settings | |
| KEYS = ["INVOICE NO", "Invoice Date", "Total Amount", "PO Number", | |
| "Payment Term", "GSTIN", "Vendor Id", "Subtotal"] | |
| REPEATS = 3 | |
| def normalise(text: str) -> str: | |
| import re | |
| import string | |
| text = text.lower() | |
| text = "".join(ch for ch in text if ch not in set(string.punctuation)) | |
| text = re.sub(r"\b(a|an|the)\b", " ", text) | |
| return " ".join(text.split()) | |
| class TorchBackend: | |
| name = "torch_fp32" | |
| def __init__(self, model_id: str, threads: int): | |
| torch.set_num_threads(threads) | |
| from transformers import pipeline | |
| self.pipe = pipeline("document-question-answering", model=model_id, | |
| device="cpu") | |
| self.tokenizer = None | |
| def answer(self, words, boxes, key): | |
| result = self.pipe(image=None, question=key, | |
| word_boxes=[[w, list(b)] for w, b in zip(words, boxes)], | |
| top_k=1) | |
| while isinstance(result, list): | |
| result = result[0] | |
| return str(result["answer"]).strip(), float(result["score"]) | |
| class OnnxBackend: | |
| def __init__(self, model_id: str, path: Path, threads: int, | |
| name: str = "onnx"): | |
| import onnxruntime as ort | |
| from transformers import AutoTokenizer | |
| self.name = name | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| options = ort.SessionOptions() | |
| options.intra_op_num_threads = threads | |
| options.inter_op_num_threads = 1 | |
| options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL | |
| self.session = ort.InferenceSession( | |
| str(path), options, providers=["CPUExecutionProvider"]) | |
| def answer(self, words, boxes, key): | |
| # text_pair must be a flat list of word strings, matching the | |
| # pipeline's LayoutLM tokenisation path. | |
| word_boxes = list(zip(words, boxes)) | |
| encoding = self.tokenizer( | |
| text=key.split(), text_pair=list(words), is_split_into_words=True, | |
| return_token_type_ids=True, truncation="only_second", | |
| max_length=self.tokenizer.model_max_length) | |
| sequence_ids = encoding.sequence_ids(0) | |
| # This must mirror DocumentQuestionAnsweringPipeline.preprocess exactly: | |
| # context token -> the word's real box | |
| # SEP token -> [1000, 1000, 1000, 1000] | |
| # everything else (question, CLS) -> [0, 0, 0, 0] | |
| # Using zeros for SEP shifts the learned position embeddings and | |
| # silently degrades answers, which is exactly what happens if this is | |
| # copied from memory instead of the pipeline source. | |
| sep_id = self.tokenizer.sep_token_id | |
| bbox = [] | |
| for token_id, sid, wid in zip(encoding["input_ids"], sequence_ids, | |
| encoding.word_ids(0)): | |
| if sid == 1: | |
| bbox.append(word_boxes[wid][1]) | |
| elif token_id == sep_id: | |
| bbox.append([1000, 1000, 1000, 1000]) | |
| else: | |
| bbox.append([0, 0, 0, 0]) | |
| feed = { | |
| "input_ids": np.array([encoding["input_ids"]], dtype=np.int64), | |
| "attention_mask": np.array([encoding["attention_mask"]], dtype=np.int64), | |
| "token_type_ids": np.array([encoding["token_type_ids"]], dtype=np.int64), | |
| "bbox": np.array([bbox], dtype=np.int64).reshape(1, -1), | |
| } | |
| start, end = self.session.run(None, feed) | |
| start, end = start[0], end[0] | |
| # Mirror the pipeline's p_mask: only context tokens may start/end a span. | |
| for i, sid in enumerate(sequence_ids): | |
| if sid != 1: | |
| start[i] = end[i] = -1e4 | |
| s = int(np.argmax(start)) | |
| e = int(np.argmax(end)) | |
| if e < s: | |
| e = s | |
| # The pipeline returns the ORIGINAL word strings joined by spaces, | |
| # not a decode of the subword ids. These differ for dates, currency | |
| # symbols and hyphenated identifiers, so they must not be conflated. | |
| word_ids = encoding.word_ids(0) | |
| w_start, w_end = word_ids[s], word_ids[e] | |
| if w_start is None or w_end is None: | |
| return "", 0.0 | |
| answer = " ".join(words[w_start:w_end + 1]) | |
| return answer, float(np.exp(start[s]) * np.exp(end[e])) | |
| def main() -> None: | |
| settings = Settings() | |
| threads = 8 | |
| model_id = settings.model_id | |
| truths = json.loads((ROOT / "corpus" / "ground_truth.json").read_text()) | |
| backends: list[object] = [TorchBackend(model_id, threads)] | |
| onnx_dir = ROOT / "artifacts" / "onnx" / model_id.replace("/", "__") | |
| if (onnx_dir / "model.onnx").exists(): | |
| backends.append(OnnxBackend(model_id, onnx_dir / "model.onnx", threads, | |
| "onnx_fp32")) | |
| if (onnx_dir / "model_int8.onnx").exists(): | |
| backends.append(OnnxBackend(model_id, onnx_dir / "model_int8.onnx", threads, | |
| "onnx_int8")) | |
| print(f"model={model_id} threads={threads} " | |
| f"documents={len(truths)} keys={len(KEYS)}\n") | |
| results: dict[str, dict] = {b.name: {"correct": 0, "total": 0, | |
| "times": [], "answers": {}} | |
| for b in backends} | |
| truth_map = {"INVOICE NO": "invoice_number", "Invoice Date": "invoice_date", | |
| "Total Amount": "total", "PO Number": "po_number", | |
| "Payment Term": "payment_terms", "GSTIN": "gstin", | |
| "Vendor Id": "vendor_tax_id", "Subtotal": "subtotal"} | |
| for record in truths: | |
| data = Path(record["path"]).read_bytes() | |
| parsed = parsing.parse(data, settings) | |
| page = parsed.pages[0] | |
| for key in KEYS: | |
| gold = str(record.get(truth_map[key], "")) | |
| for backend in backends: | |
| backend.answer(page.words, page.boxes, key) # warm | |
| started = time.perf_counter() | |
| answer, _score = backend.answer(page.words, page.boxes, key) | |
| elapsed = time.perf_counter() - started | |
| slot = results[backend.name] | |
| slot["times"].append(elapsed) | |
| slot["answers"][(record["document_id"], key)] = answer | |
| if normalise(answer) == normalise(gold): | |
| slot["correct"] += 1 | |
| slot["total"] += 1 | |
| print("=" * 74) | |
| print(f"{'backend':<14} {'ms/key':>9} {'speedup':>9} {'accuracy':>10} " | |
| f"{'agreement w/ torch':>21}") | |
| print("=" * 74) | |
| baseline_ms = statistics.median(results["torch_fp32"]["times"]) * 1000 | |
| torch_answers = results["torch_fp32"]["answers"] | |
| for backend in backends: | |
| slot = results[backend.name] | |
| ms = statistics.median(slot["times"]) * 1000 | |
| acc = slot["correct"] / slot["total"] | |
| agree = sum(1 for k, v in slot["answers"].items() | |
| if normalise(v) == normalise(torch_answers.get(k, ""))) | |
| total_keys = len(slot["answers"]) | |
| agree_str = f"{agree}/{total_keys}" | |
| print(f"{backend.name:<14} {ms:>9.0f} {baseline_ms/ms:>8.2f}x " | |
| f"{acc:>9.1%} {agree_str:>21}") | |
| print("\nPER-KEY ACCURACY") | |
| keys_hdr = f"{'backend':<14}" + "".join(f"{k[:13]:>15}" for k in KEYS) | |
| print(keys_hdr) | |
| for backend in backends: | |
| slot = results[backend.name] | |
| row = f"{backend.name:<14}" | |
| for key in KEYS: | |
| hits = total = 0 | |
| for record in truths: | |
| gold = str(record.get(truth_map[key], "")) | |
| pred = slot["answers"].get((record["document_id"], key), "") | |
| total += 1 | |
| hits += normalise(pred) == normalise(gold) | |
| row += f"{hits/total:>14.0%} " | |
| print(row) | |
| json.dump({name: {"median_ms": statistics.median(slot["times"]) * 1000, | |
| "accuracy": slot["correct"] / slot["total"], | |
| "n": slot["total"]} | |
| for name, slot in results.items()}, | |
| (ROOT / "results" / "backend_comparison.json").open("w"), | |
| indent=2) | |
| if __name__ == "__main__": | |
| main() | |