File size: 5,849 Bytes
7b2177e | 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 | # GENERATED by hub/build_hub_package.py from src/bodhan_genai/ocr/layout/modeling_ppdoc.py -- do not edit.
# Vendored so this repo is self-contained: `pip install transformers torch pillow` is the
# whole install. See indic_doc_parser.py for usage.
"""IndicDocLayout: trainable PP-DocLayoutV3 (document-pretrained strong init).
HF ships PPDocLayoutV3ForObjectDetection inference-only (forward raises on labels).
This subclass unblocks training: it calls the inner model with labels (which builds the
contrastive-denoising groups), reuses the base RT-DETR detection loss on its outputs, and
adds our locality-weighted GCE order loss on its (pretrained) order_logits.
Backbone + decoder + order/mask heads start from the document-pretrained checkpoint;
only the class heads are re-init'd for our 37 education classes.
"""
from dataclasses import dataclass
import torch
import torch.nn as nn
from transformers import PPDocLayoutV3Config, PPDocLayoutV3ForObjectDetection
from transformers.loss.loss_rt_detr import RTDetrHungarianMatcher
from transformers.utils import ModelOutput
from idp_model_order_loss import locality_gce
@dataclass
class PPDocOutput(ModelOutput):
loss: torch.FloatTensor | None = None
logits: torch.FloatTensor | None = None
pred_boxes: torch.FloatTensor | None = None
order_logits: torch.FloatTensor | None = None
last_hidden_state: torch.FloatTensor | None = None
class PPDocLayoutV3Trainable(PPDocLayoutV3ForObjectDetection):
def __init__(self, config):
super().__init__(config)
self.lambda_order = getattr(config, "lambda_order", 5.0)
# Built on first use, not here: the matcher is training-only, and constructing it calls
# requires_backends(["scipy"]). Eagerly, that makes scipy a hard dependency of merely
# LOADING the detector -- so inference-only installs fail on import with a library they
# will never call. It holds no parameters, so this does not change the state dict.
self._matcher = None
self.loss_type = "RTDetrForObjectDetection" # base RT-DETR loss over its outputs
@classmethod
def build(cls, ckpt, num_labels, id2label, label2id, lambda_order=5.0):
config = PPDocLayoutV3Config.from_pretrained(
ckpt, num_labels=num_labels, id2label=id2label, label2id=label2id
)
config.lambda_order = lambda_order
config.loss_type = "RTDetrForObjectDetection"
# PP-DocLayoutV3's denoising path is buggy (embed size num_labels but pads with
# num_labels -> index error); it was never run since HF blocks training. Disable it
# (optional convergence aid). Re-enable later by resizing denoising_class_embed to +1.
config.num_denoising = 0
# RT-DETR loss/matcher fields the base config lacks
defaults = {
"use_focal_loss": True,
"auxiliary_loss": True,
"weight_loss_vfl": 1.0,
"weight_loss_bbox": 5.0,
"weight_loss_giou": 2.0,
"matcher_class_cost": 2.0,
"matcher_bbox_cost": 5.0,
"matcher_giou_cost": 2.0,
"matcher_alpha": 0.25,
"matcher_gamma": 2.0,
"focal_loss_alpha": 0.25,
"focal_loss_gamma": 2.0,
"eos_coefficient": 1e-4,
}
for k, v in defaults.items():
if not hasattr(config, k):
setattr(config, k, v)
model = cls.from_pretrained(ckpt, config=config, ignore_mismatched_sizes=True)
# re-init class heads (paddle doc classes -> our education classes); keep everything else
for m in model.modules():
if (isinstance(m, nn.Linear) and m.out_features == num_labels) or (
isinstance(m, nn.Embedding) and m.num_embeddings == num_labels + 1
):
m.reset_parameters()
return model
def _order_loss(self, order_logits, logits, pred_boxes, labels):
if self._matcher is None:
self._matcher = RTDetrHungarianMatcher(self.config) # needs scipy; training only
idx = self._matcher({"logits": logits, "pred_boxes": pred_boxes}, labels)
tot, n = 0.0, 0
for b, (src, tgt) in enumerate(idx):
if src.numel() < 2:
continue
order = labels[b]["reading_order"][tgt]
S = order_logits[b][src][:, src]
tot = tot + locality_gce(S, order)
n += 1
return tot / max(n, 1) if n else order_logits.sum() * 0.0
def forward(self, pixel_values, pixel_mask=None, labels=None, **kwargs):
outputs = self.model(pixel_values=pixel_values, pixel_mask=pixel_mask, labels=labels)
dn = outputs.denoising_meta_values if self.training else None
outputs_class = outputs.intermediate_logits
outputs_coord = outputs.intermediate_reference_points
logits, pred_boxes = outputs_class[:, -1], outputs_coord[:, -1]
order_logits = outputs.out_order_logits[:, -1] # [B, num_queries, num_queries]
loss = None
if labels is not None:
loss, _, _ = self.loss_function(
logits,
labels,
self.device,
pred_boxes,
self.config,
outputs_class,
outputs_coord,
enc_topk_logits=outputs.enc_topk_logits,
enc_topk_bboxes=outputs.enc_topk_bboxes,
denoising_meta_values=dn,
)
loss = loss + self.lambda_order * self._order_loss(
order_logits, logits, pred_boxes, labels
)
return PPDocOutput(
loss=loss,
logits=logits,
pred_boxes=pred_boxes,
order_logits=order_logits,
last_hidden_state=outputs.last_hidden_state,
)
|