File size: 4,950 Bytes
44c2f50
 
 
 
 
 
 
 
 
 
6267e20
 
44c2f50
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""CORD dataset loader + normalizer.

CORD (Consolidated Receipt Dataset, Naver Clova) has ~1000 receipts with rich
ground truth including line items, subtotal, tax, and total — much richer
than SROIE.

Reference: https://github.com/clovaai/cord
"""
from __future__ import annotations

from collections.abc import Iterator
from typing import Any

from src.data_prep.parsers import clean_text, parse_money
from src.schemas import Receipt, ReceiptLineItem
from src.utils.logging import logger

DEFAULT_DATASET_IDS = (
    "naver-clova-ix/cord-v2",
    "katanaml-org/invoices-donut-data-v1",  # occasional CORD mirror
)


def load_cord_split(split: str = "test", dataset_id: str | None = None):
    """Load a CORD split from Hugging Face."""
    from datasets import load_dataset

    ids_to_try = (dataset_id,) if dataset_id else DEFAULT_DATASET_IDS
    last_err: Exception | None = None

    for ds_id in ids_to_try:
        try:
            logger.info(f"Loading CORD split={split!r} from Hugging Face id={ds_id!r}")
            return load_dataset(ds_id, split=split)
        except Exception as e:
            logger.warning(f"Failed to load {ds_id}: {e}")
            last_err = e

    raise RuntimeError(
        f"Could not load any CORD dataset. Last error: {last_err}"
    ) from last_err


def _extract_menu_items(menu: list[dict[str, Any]]) -> list[ReceiptLineItem]:
    """Parse CORD's `menu` list into ReceiptLineItem objects.

    CORD menu items look like:
      {"nm": "product name", "cnt": "1", "price": "5.00", "unitprice": "5.00"}
    """
    items: list[ReceiptLineItem] = []
    for m in menu:
        description = clean_text(m.get("nm"))
        if not description:
            continue
        try:
            qty_raw = m.get("cnt")
            qty = float(qty_raw) if qty_raw not in (None, "") else None
        except (TypeError, ValueError):
            qty = None

        items.append(
            ReceiptLineItem(
                description=description,
                quantity=qty,
                unit_price=parse_money(m.get("unitprice") or m.get("price")),
                total=parse_money(m.get("price")),
            )
        )
    return items


def normalize_cord_ground_truth(gt: dict[str, Any]) -> Receipt | None:
    """Normalize a CORD ground_truth JSON block into our Receipt schema.

    CORD `gt_parse` structure (top-level keys):
      {
        "menu": [ {...}, ... ],
        "sub_total": { "subtotal_price": "...", "tax_price": "...", ... },
        "total": { "total_price": "...", "cashprice": "...", ... }
      }
    """
    # CORD often nests everything under "gt_parse" or similar; unwrap if needed.
    src = gt.get("gt_parse") or gt.get("valid_line") or gt

    menu = src.get("menu")
    if isinstance(menu, dict):
        # Some CORD variants have a single menu dict instead of a list.
        menu = [menu]
    if not isinstance(menu, list):
        menu = []

    line_items = _extract_menu_items(menu)

    sub_total_block = src.get("sub_total") or {}
    total_block = src.get("total") or {}

    subtotal = parse_money(sub_total_block.get("subtotal_price"))
    tax = parse_money(sub_total_block.get("tax_price"))
    total = parse_money(total_block.get("total_price") or total_block.get("cashprice"))

    if total is None:
        logger.debug(f"CORD record missing total; keys={list(src.keys())}")
        return None

    # CORD receipts often don't include merchant name in the structured GT.
    # We fall back to "Unknown merchant" — evaluation still works on the fields
    # that CORD does provide (line items, subtotal, tax, total).
    merchant = clean_text(src.get("merchant") or src.get("nm")) or "Unknown merchant"

    try:
        return Receipt(
            merchant=merchant,
            line_items=line_items,
            subtotal=subtotal,
            tax=tax,
            total=total,
            # CORD is Korean receipts by origin. Use KRW as default.
            currency="KRW",
        )
    except Exception as e:
        logger.warning(f"Failed to build Receipt from CORD record: {e}")
        return None


def iter_normalized(dataset) -> Iterator[tuple[str, Receipt]]:
    """Iterate a HF CORD Dataset, yielding (record_id, Receipt) pairs."""
    import json as _json

    for idx, rec in enumerate(dataset):
        record_id = str(rec.get("id") or f"cord_{idx:05d}")
        # CORD ground truth is usually stored as a JSON string under `ground_truth`.
        gt_raw = rec.get("ground_truth") or rec.get("gt_parse") or rec
        if isinstance(gt_raw, str):
            try:
                gt = _json.loads(gt_raw)
            except _json.JSONDecodeError:
                logger.debug(f"Could not JSON-decode CORD ground_truth for record {record_id}")
                continue
        else:
            gt = gt_raw

        normalized = normalize_cord_ground_truth(gt)
        if normalized is not None:
            yield record_id, normalized