question_router.py

#!/usr/bin/env python3
"""Rule-based question router.

The reader network never sees the question. This module turns a free-form
Korean or English question into a field selector, then indexes the record the
network produced. It is dependency-free and is shipped alongside the model,
because the ONNX graph alone cannot answer a question.

The model reads exactly two things: the phone number's digits and the street
number's digits. Anything else must return an empty answer. A router that
guesses is worse than one that declines, because a plausible wrong number is
indistinguishable from a right one downstream.

Divergence from `tiny_receipt_vqa/train.py`
-------------------------------------------
The regex bodies started as a verbatim copy of that file's
`route_family_from_question` and `phone_op_from_question`. They have since been
corrected, because in the baseline the router only picked an adapter — a wrong
route still produced an answer from the network — whereas here the router
*decides the answer*. The same code is far more dangerous in this position.

Four fixes, each covered by `test_question_router.py`:

1. Item, price, and store-name questions are matched *before* address. The
   Korean particle `로` is a substring of ordinary words (`합계로`, `제품으로`,
   `세로`), so an address test that runs first swallows them.
2. An address question only yields the street number when it actually asks for
   a number. `가게 주소가 무엇입니까?` now returns nothing instead of the
   street number.
3. Digit indices are range-checked. `digit 0` used to index `phone[-1]` and
   return the last digit.
4. `last digit` / `끝자리` with no explicit ordinal now resolve to `back_1`
   instead of falling through to an empty answer.
"""
from __future__ import annotations

import re
import unicodedata

__all__ = [
    "clean_text",
    "digits_only",
    "normalize_address_text",
    "street_no_from_address",
    "ordinal_en",
    "route_family_from_question",
    "phone_op_from_question",
    "address_op_from_question",
    "answer_from_record",
    "SUPPORTED_FAMILIES",
]

SUPPORTED_FAMILIES = ("phone", "address")

_ORDINALS = {"first": 1, "second": 2, "third": 3, "fourth": 4,
             "fifth": 5, "sixth": 6, "seventh": 7, "eighth": 8,
             "ninth": 9, "tenth": 10}


def clean_text(s: object) -> str:
    s = str(s if s is not None else "")
    return unicodedata.normalize("NFC", re.sub(r"\s+", " ", s.replace("\n", " ")).strip())


def digits_only(s: object) -> str:
    return "".join(re.findall(r"\d", str(s if s is not None else "")))


def normalize_address_text(s: object) -> str:
    text = clean_text(s)
    return re.sub(r"(\d+(?:\s+\d+)+)$", lambda m: re.sub(r"\s+", "", m.group(1)), text)


def street_no_from_address(address: object) -> str:
    text = normalize_address_text(address)
    match = re.search(r"(\d+(?:\s+\d+)*)\s*$", text)
    return digits_only(match.group(1)) if match else ""


def ordinal_en(n: int) -> str:
    names = {v: k for k, v in _ORDINALS.items()}
    return names.get(n, f"{n}th")


def _is_phone(q: str, lower: str) -> bool:
    return bool(re.search(r"\b(?:phone|telephone|tel)\b", lower)) or \
        any(k in q for k in ("전화", "폰번호", "연락처"))


def _is_item(q: str, lower: str) -> bool:
    return bool(re.search(r"\b(?:item|items|product|products|price|prices|qty|quantity|"
                          r"total|totals|subtotal|amount|cost|purchased|bought)\b", lower)) or \
        any(k in q for k in ("품목", "상품", "제품", "가격", "단가", "수량",
                             "구매", "구입", "총액", "금액", "합계", "개수"))


def _is_store_name(q: str, lower: str) -> bool:
    return bool(re.search(r"\b(?:store name|shop name|merchant|business name|"
                          r"name of the (?:store|shop))\b", lower)) or \
        any(k in q for k in ("상호", "가게 이름", "매장 이름", "점포 이름", "가게명", "매장명"))


def _is_address(q: str, lower: str) -> bool:
    return bool(re.search(r"\b(?:address|street|road|location)\b", lower)) or \
        bool(re.search(r"\b(?:st|rd)\.", lower)) or \
        any(k in q for k in ("주소", "도로명", "위치", "번지")) or \
        bool(re.search(r"[가-힣]{1,10}(?:길|로|대로)\s*\[?\?", q))


def route_family_from_question(question: object) -> str:
    """'phone', 'address', 'item', 'store', or 'other'.

    Only 'phone' and 'address' are answerable; the rest exist so that an
    unsupported question is recognised rather than mistaken for a supported one.
    Order matters — see the module docstring.
    """
    q = clean_text(question)
    lower = q.lower()
    if _is_phone(q, lower):
        return "phone"
    if _is_item(q, lower):
        return "item"
    if _is_store_name(q, lower):
        return "store"
    if _is_address(q, lower):
        return "address"
    return "other"


def phone_op_from_question(question: str) -> str:
    """'front_N', 'back_N', or 'phone_digit' when no index can be recovered."""
    q = clean_text(question).lower()
    op = "phone_digit"
    m = re.search(r"\b(" + "|".join(_ORDINALS) + r")\b", q)
    if m:
        op = f"front_{_ORDINALS[m.group(1)]}"
    m = re.search(r"(?:front of|from the front|digit)\D*(\d+)", q)
    if m:
        op = f"front_{m.group(1)}"
    m = re.search(r"(?:from the end|from the back|from the right|last)\D*(\d+)", q)
    if m:
        op = f"back_{m.group(1)}"
    m = re.search(r"앞에서\s*(\d+)\s*번째", q)
    if m:
        op = f"front_{m.group(1)}"
    m = re.search(r"(?:앞|앞자리)\s*(\d+)\s*(?:번째|번)?", q)
    if m:
        op = f"front_{m.group(1)}"
    m = re.search(r"뒤에서\s*(\d+)\s*번째", q)
    if m:
        op = f"back_{m.group(1)}"
    m = re.search(r"(?:뒤|뒷자리|끝자리)\s*(\d+)\s*(?:번째|번)?", q)
    if m:
        op = f"back_{m.group(1)}"
    if any(k in q for k in ("from the back", "from the end", "from last")):
        m = re.search(r"\b(" + "|".join(_ORDINALS) + r")\b", q)
        if m:
            op = f"back_{_ORDINALS[m.group(1)]}"
    if op == "phone_digit":
        # "last digit", "끝자리", "마지막 숫자" carry an ordinal of one.
        if re.search(r"\blast\b", q) or any(k in q for k in ("끝자리", "마지막")):
            op = "back_1"
        elif re.search(r"\bfirst\b", q) or "첫자리" in q or "첫 번째" in q:
            op = "front_1"
    return op


def address_op_from_question(question: object) -> str:
    """'street_no' when the question asks for the street number, else 'unsupported'.

    The model cannot transcribe an address, so a question about the address text
    has to decline rather than hand back the number it happens to hold.
    """
    q = clean_text(question)
    lower = q.lower()
    if "[?]" in q or "?]" in q:
        return "street_no"
    if any(k in lower for k in ("fill the blank", "fill in the blank")):
        return "street_no"
    if any(k in q for k in ("빈 칸", "빈칸")):
        return "street_no"
    if re.search(r"\b(?:street|road|building|house|block)\s*(?:number|no\.?|num)\b", lower):
        return "street_no"
    if re.search(r"\bnumber\b.*\b(?:address|street|road)\b", lower) or \
            re.search(r"\b(?:address|street|road)\b.*\bnumber\b", lower):
        return "street_no"
    if "번지" in q:
        return "street_no"
    if re.search(r"(?:도로명|주소|위치|길|로)\s*(?:뒤|뒤의|끝|마지막)?\s*(?:에)?\s*숫자", q):
        return "street_no"
    if re.search(r"숫자", q) and any(k in q for k in ("주소", "도로명", "위치")):
        return "street_no"
    return "unsupported"


def answer_from_record(question: str, phone: str, street: str) -> str:
    """Index an already-read record. Returns '' when the question cannot be served."""
    family = route_family_from_question(question)
    if family == "address":
        return street if address_op_from_question(question) == "street_no" else ""
    if family == "phone":
        m = re.match(r"(front|back)_(\d+)$", phone_op_from_question(question))
        if not m or not phone:
            return ""
        i = int(m.group(2))
        if not 1 <= i <= len(phone):
            return ""
        return phone[i - 1] if m.group(1) == "front" else phone[-i]
    return ""