LayoutLMv3 KIE Pipeline

Pipeline fine-tune LayoutLMv3 (hoặc LayoutXLM) cho bài toán Key Information Extraction (KIE) — trích xuất thông tin từ hóa đơn, biên lai, và các loại tài liệu bán cấu trúc khác.

Thiết kế xoay quanh 1 nguyên tắc cốt lõi: đổi dataset, đổi ngôn ngữ, hay đổi kiến trúc model chỉ cần sửa file YAML — không bao giờ phải đụng vào code Python.

• Config-driven — toàn bộ field của dataset, ngôn ngữ OCR, kiến trúc model, và m��i hyperparameter đều nằm trong 1 file .yaml duy nhất.
• Generic dataset loader — chạy được với bất kỳ dataset nào theo đúng cấu trúc thư mục img/box/entities, không hard-code riêng cho 1 dataset cụ thể.
• Hỗ trợ nhiều kiến trúc — đổi qua lại giữa LayoutLMv3 (tiếng Anh) và LayoutXLM (đa ngôn ngữ, phù hợp tiếng Việt) chỉ bằng 1 dòng config.
• An toàn về RAM — data được nạp qua generator (Dataset.from_generator) thay vì load hết vào RAM cùng lúc, tránh bị hệ điều hành kill tiến trình (OOM) với dataset lớn.
• Resume training — tự động tìm và tiếp tục từ checkpoint hợp lệ gần nhất, tự bỏ qua checkpoint bị hỏng do ngắt giữa chừng (Ctrl+C đúng lúc đang ghi file).
• Checkpoint tự chứa đầy đủ — mỗi checkpoint tự động lưu kèm processor cùng với trọng số model, nên checkpoint nào cũng load & test độc lập được ngay (giống cách YOLO lưu epochN.pt/best.pt).
• Logging kép — mọi lệnh chạy đều ghi song song ra file log (đã lọc sạch nhiễu progress bar) và lưu toàn bộ lịch sử metric theo từng step ra JSON để vẽ biểu đồ sau này.
• Đánh giá theo từng entity — precision/recall/F1 tách riêng theo từng loại field, không chỉ 1 con số F1 tổng.
• Nhất quán giữa train và inference — logic tách word dùng để gán nhãn lúc train được tái sử dụng y hệt lúc predict, tránh lệch giữa training và thực tế sử dụng (train/serve skew).

Tài liệu gốc (img/box/entities)
        │
        ▼
┌───────────────────┐
│   data_loader.py    │  đọc file box OCR, tách từng dòng thành từng word riêng,
│                     │  gán nhãn BIO bằng cách so khớp entity value
└─────────┬──────────┘
          ▼
┌───────────────────┐
│  preprocessing.py   │  sinh label list, tokenize + align nhãn theo subword
└─────────┬──────────┘
          ▼
┌───────────────────┐
│   model_utils.py     │  registry kiến trúc (LayoutLMv3 / LayoutXLM),
│                     │  metric, class-weighting, callback lưu checkpoint
└─────────┬──────────┘
          ▼
   train.py / eval.py / predict.py

layoutlmv3_kie/
├── configs/
│   ├── sroie.yaml          # hóa đơn tiếng Anh (SROIE)
│   └── mcocr_vi.yaml       # hóa đơn tiếng Việt (MC-OCR)
├── data_loader.py           # loader chung: đọc box OCR, tách word, gán nhãn BIO
├── preprocessing.py          # sinh label list, tokenize, gộp entity
├── model_utils.py             # load config, registry kiến trúc, metric, logging, callback
├── train.py                   # script train (config-driven, resume được)
├── eval.py                    # đánh giá chi tiết theo từng entity, chạy trên checkpoint bất kỳ
├── predict.py                  # OCR + inference 1 ảnh (CLI + class tái sử dụng)
├── requirements.txt
└── README.md

git clone <repo-cua-ban>
cd layoutlmv3_kie
pip install -r requirements.txt

Lưu ý (Linux/WSL2): cài paddlepaddle-gpu có thể cần thêm --break-system-packages tùy môi trường. Nếu gặp lỗi CUDA/GPU, chuyển sang bản CPU của paddlepaddle.

Mọi dataset đều phải theo đúng cấu trúc:

data/<ten_dataset>/
├── train/
│   ├── img/         *.jpg
│   ├── box/          *.txt   (mỗi dòng 1 vùng OCR: x0,y0,x1,y1,x2,y2,x3,y3,text)
│   └── entities/     *.txt   (JSON phẳng: {"company": "...", "date": "...", ...})
└── test/
    └── (cấu trúc tương tự)

• box/*.txt — mỗi dòng là 1 tứ giác (4 điểm góc) kèm text OCR nhận diện được trong vùng đó. 1 dòng có thể chứa nhiều từ (ví dụ nguyên cả tên công ty) — loader tự động tách thành từng word riêng, các word cùng dòng dùng chung 1 bbox.
• entities/*.txt — JSON phẳng chứa giá trị ground-truth của từng field cần trích xuất.

Nhãn BIO được sinh ra bằng cách so khớp chính xác (exact string match) giữa entity value và các word liên tiếp trong kết quả OCR. Nếu tỷ lệ match thấp, kiểm tra lại sự không nhất quán giữa OCR và annotation (xem mục Giới hạn đã biết).

Copy configs/sroie.yaml, sửa các field phù hợp với dataset của bạn (xem Tài liệu Config).

python train.py --config configs/sroie.yaml

Script sẽ in ra tỷ lệ match entity trên tập train/test trước khi bắt đầu train — kiểm tra tỷ lệ này đủ cao (>85%) trước khi để chạy hết.

python train.py --config configs/sroie.yaml --resume

Tự động tìm checkpoint hợp lệ gần nhất (bỏ qua checkpoint bị hỏng do ngắt giữa lúc lưu), resume đúng optimizer state, learning rate schedule, và epoch hiện tại.

python eval.py --config configs/sroie.yaml --checkpoint outputs/sroie_run1/final

python predict.py --config configs/sroie.yaml \
    --checkpoint outputs/sroie_run1/final \
    --image invoice.jpg

KIEPredictor trong predict.py cũng có thể import trực tiếp để dùng trong service (load model 1 lần, gọi .predict() nhiều lần):

from predict import KIEPredictor

predictor = KIEPredictor("configs/sroie.yaml", "outputs/sroie_run1/final")
result = predictor.predict("invoice.jpg")
# {'COMPANY': ['ABC TRADING'], 'DATE': ['15/01/2019'], 'TOTAL': ['193.00']}

dataset:
  root_dir: ./data/SROIE2019
  ocr_lang: en                      # mã ngôn ngữ PaddleOCR, dùng trong predict.py
  cache_dir: null                   # tùy chọn; mặc định là <root_dir>/.hf_cache/<split>

  entity_fields:                    # field JSON trong entities/*.txt -> tên nhãn BIO
    - json_key: company
      label: COMPANY
    - json_key: date
      label: DATE
    - json_key: address
      label: ADDRESS
    - json_key: total
      label: TOTAL

model:
  architecture: layoutlmv3          # hoặc: layoutxlm
  base_checkpoint: microsoft/layoutlmv3-base
  max_length: 512

training:
  output_dir: ./outputs/sroie_run1
  batch_size: 2
  eval_batch_size: 2
  grad_accum_steps: 8
  learning_rate: 2e-5
  num_epochs: 50
  warmup_ratio: 0.1
  weight_decay: 0.01
  max_grad_norm: 1.0
  fp16: true
  logging_steps: 10
  save_total_limit: 2
  early_stopping_patience: 5
  use_class_weights: false          # bật nếu entity bị mất cân bằng nặng
  dataloader_num_workers: 2
  seed: 42

1. Chuẩn bị data đúng format ở trên.
2. Copy configs/sroie.yaml → configs/<ten>.yaml, sửa phần dataset.*.
3. Chạy với --config configs/<ten>.yaml. Không cần sửa bất kỳ dòng code Python nào.

model:
  architecture: layoutxlm
  base_checkpoint: microsoft/layoutxlm-base

Registry kiến trúc nằm trong model_utils.py (dict ARCHITECTURES) — thêm kiến trúc khác (vd: Donut, LiLT) chỉ cần thêm 1 entry vào đó, không cần sửa train.py, eval.py, hay predict.py.

Output console được ghi song song ra file, đã lọc sạch nhiễu progress bar. Xóa <root_dir>/.hf_cache để ép parse lại data từ đầu (vd sau khi đổi entity_fields).

Pipeline này dùng so khớp chuỗi chính xác (exact string match) để sinh nhãn BIO từ cặp (OCR words, entity JSON). Cách này đơn giản, minh bạch, nhưng có vài giới hạn thật cần biết:

• Nhiễu trong annotation: nếu giá trị entity ground-truth không khớp chính xác với text OCR (lỗi chính tả, lỗi OCR, format không nhất quán), sample đó sẽ không được gán nhãn. Với SROIE, điều này ảnh hưởng 1 tỷ lệ nhỏ sample (ví dụ: annotation và OCR lệch nhau 1 chữ số/chữ cái).
• Entity nhiều từ cần tokenize đúng ở cấp độ word: OCR engine thường detect nguyên cả 1 dòng (ví dụ cả tên công ty, hay "Date 02/03/2018") thành 1 vùng text duy nhất. Pipeline này tách mỗi vùng detect được thành từng word riêng ở cả lúc train lẫn lúc inference, dùng chung 1 logic — sự nhất quán này bắt buộc phải có; nếu tách lúc train nhưng không tách lúc inference (hoặc ngược lại), model sẽ nhận input ở granularity khác hẳn những gì nó từng học, và âm thầm cho kết quả kém đi.
• Không hiểu cấu trúc bảng biểu, không xử lý tài liệu nhiều trang, và bị cắt bởi max_length với tài liệu dài hơn giới hạn token đã cấu hình.

Với các loại tài liệu mà những giới hạn này gây ảnh hưởng lớn (hợp đồng dài, bảng biểu phức tạp, scan nhiều trang), hướng token-classification như pipeline này không phải lựa chọn phù hợp — nên cân nhắc hướng VLM-based hoặc document-QA thay thế.

Xem requirements.txt. Các dependency chính: transformers, datasets, evaluate, seqeval, accelerate, torch, paddleocr, pyyaml, sentencepiece (bắt buộc cho tokenizer của LayoutXLM).

Thêm license của bạn ở đây (vd: MIT). Lưu ý microsoft/layoutlmv3-base và microsoft/layoutxlm-base được Microsoft phân phối theo giấy phép CC-BY-NC-SA-4.0 — chỉ dùng cho mục đích phi thương mại trừ khi có thỏa thuận riêng với Microsoft. Kiểm tra license của checkpoint gốc trước khi dùng cho mục đích thương mại.