Spaces:
Running
Running
| """Align corrected words back to the original TrOCR words' confidence. | |
| Ported from post_correction_full_test_v2.ipynb (`align_words`/`_normalise`). Page-level | |
| correction rewrites text, so to keep the per-word confidence tint in the Reader we map | |
| each corrected word to the original word(s) it came from via difflib.SequenceMatcher: | |
| equal -> keep the original word's confidence | |
| replace -> average confidence of the replaced originals, flagged qwen_replaced | |
| insert -> a word Qwen added with no original -> default confidence, qwen_replaced | |
| delete -> original dropped -> skipped | |
| Returns a fresh list[Word] for the corrected line. | |
| """ | |
| from __future__ import annotations | |
| import difflib | |
| import re | |
| from ..schemas import Word | |
| # Confidence for words Qwen inserted that have no original counterpart. | |
| DEFAULT_INSERT_CONF = 0.5 | |
| def _normalise(word: str) -> str: | |
| return re.sub(r"[^\w]", "", word.lower()) | |
| def align_line( | |
| original_words: list[Word], | |
| corrected_text: str, | |
| low_conf_threshold: float = 0.60, | |
| ) -> list[Word]: | |
| """Map the corrected line's words onto the original words' confidences.""" | |
| corr_tokens = corrected_text.split() | |
| if not corr_tokens: | |
| return [] | |
| if not original_words: | |
| return [ | |
| Word(text=t, confidence=DEFAULT_INSERT_CONF, | |
| needs_review=DEFAULT_INSERT_CONF < low_conf_threshold, qwen_replaced=True) | |
| for t in corr_tokens | |
| ] | |
| orig_keys = [_normalise(w.text) for w in original_words] | |
| corr_keys = [_normalise(t) for t in corr_tokens] | |
| matcher = difflib.SequenceMatcher(None, orig_keys, corr_keys, autojunk=False) | |
| out: list[Word] = [] | |
| for tag, i1, i2, j1, j2 in matcher.get_opcodes(): | |
| if tag == "equal": | |
| for i, j in zip(range(i1, i2), range(j1, j2)): | |
| o = original_words[i] | |
| out.append(Word(text=corr_tokens[j], confidence=o.confidence, | |
| needs_review=o.needs_review, qwen_replaced=False)) | |
| elif tag == "replace": | |
| confs = [original_words[i].confidence for i in range(i1, i2)] | |
| avg = sum(confs) / len(confs) if confs else DEFAULT_INSERT_CONF | |
| for j in range(j1, j2): | |
| out.append(Word(text=corr_tokens[j], confidence=round(avg, 4), | |
| needs_review=avg < low_conf_threshold, qwen_replaced=True)) | |
| elif tag == "insert": | |
| for j in range(j1, j2): | |
| out.append(Word(text=corr_tokens[j], confidence=DEFAULT_INSERT_CONF, | |
| needs_review=DEFAULT_INSERT_CONF < low_conf_threshold, | |
| qwen_replaced=True)) | |
| # tag == "delete": original word dropped by the corrector -> nothing emitted | |
| return out | |
| def align_page( | |
| conf_words: list[tuple[str, float]], | |
| corrected_lines: list[str], | |
| low_conf_threshold: float = 0.60, | |
| ) -> list[list[Word]]: | |
| """Page-level alignment (faithful to post_correction_full_test_v2.ipynb). | |
| Aligns the flat stream of confidence-scored raw words against the flat stream | |
| of corrected words (which carry their corrected-line index), then groups the | |
| resulting words back by corrected line. Robust when the corrected text has a | |
| different line structure than the raw OCR (common here). | |
| Opcodes: equal -> original confidence; replace -> avg confidence of the | |
| replaced originals + qwen_replaced; insert -> default confidence (grey, not | |
| flagged, matching the v2 HTML); delete -> dropped. | |
| Returns one list[Word] per corrected line. | |
| """ | |
| n_lines = len(corrected_lines) | |
| grouped: list[list[Word]] = [[] for _ in range(n_lines)] | |
| # flat corrected words tagged with their line index | |
| corr_wl: list[tuple[int, str]] = [] | |
| for li, line in enumerate(corrected_lines): | |
| for w in line.split(): | |
| corr_wl.append((li, w)) | |
| if not corr_wl: | |
| return grouped | |
| orig_keys = [_normalise(w) for w, _ in conf_words] | |
| corr_keys = [_normalise(w) for _, w in corr_wl] | |
| matcher = difflib.SequenceMatcher(None, orig_keys, corr_keys, autojunk=False) | |
| def emit(li: int, text: str, conf: float, replaced: bool) -> None: | |
| grouped[li].append(Word(text=text, confidence=round(conf, 4), | |
| needs_review=conf < low_conf_threshold, | |
| qwen_replaced=replaced)) | |
| for tag, i1, i2, j1, j2 in matcher.get_opcodes(): | |
| if tag == "equal": | |
| for i, j in zip(range(i1, i2), range(j1, j2)): | |
| li, w = corr_wl[j] | |
| emit(li, w, conf_words[i][1], False) | |
| elif tag == "replace": | |
| confs = [conf_words[i][1] for i in range(i1, i2)] | |
| avg = sum(confs) / len(confs) if confs else DEFAULT_INSERT_CONF | |
| for j in range(j1, j2): | |
| li, w = corr_wl[j] | |
| emit(li, w, avg, True) | |
| elif tag == "insert": | |
| for j in range(j1, j2): | |
| li, w = corr_wl[j] | |
| emit(li, w, DEFAULT_INSERT_CONF, False) | |
| # delete: raw word dropped -> nothing | |
| return grouped | |