|
""" |
|
Donut |
|
Copyright (c) 2022-present NAVER Corp. |
|
MIT License |
|
""" |
|
import json |
|
import os |
|
import random |
|
from collections import defaultdict |
|
from typing import Any, Dict, List, Tuple, Union |
|
|
|
import torch |
|
import zss |
|
from datasets import load_dataset |
|
from nltk import edit_distance |
|
from torch.utils.data import Dataset |
|
from transformers.modeling_utils import PreTrainedModel |
|
from zss import Node |
|
|
|
|
|
def save_json(write_path: Union[str, bytes, os.PathLike], save_obj: Any): |
|
with open(write_path, "w") as f: |
|
json.dump(save_obj, f) |
|
|
|
|
|
def load_json(json_path: Union[str, bytes, os.PathLike]): |
|
with open(json_path, "r") as f: |
|
return json.load(f) |
|
|
|
|
|
class DonutDataset(Dataset): |
|
""" |
|
DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets) |
|
Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt), |
|
and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string) |
|
|
|
Args: |
|
dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl |
|
ignore_id: ignore_index for torch.nn.CrossEntropyLoss |
|
task_start_token: the special token to be fed to the decoder to conduct the target task |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dataset_name_or_path: str, |
|
donut_model: PreTrainedModel, |
|
max_length: int, |
|
split: str = "train", |
|
ignore_id: int = -100, |
|
task_start_token: str = "<s>", |
|
prompt_end_token: str = None, |
|
sort_json_key: bool = True, |
|
): |
|
super().__init__() |
|
|
|
self.donut_model = donut_model |
|
self.max_length = max_length |
|
self.split = split |
|
self.ignore_id = ignore_id |
|
self.task_start_token = task_start_token |
|
self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token |
|
self.sort_json_key = sort_json_key |
|
|
|
self.dataset = load_dataset(dataset_name_or_path, split=self.split) |
|
self.dataset_length = len(self.dataset) |
|
|
|
self.gt_token_sequences = [] |
|
|
|
for sample in self.dataset: |
|
|
|
|
|
ground_truth = json.loads(sample["ground_truth"]) |
|
|
|
if "gt_parses" in ground_truth: |
|
assert isinstance(ground_truth["gt_parses"], list) |
|
gt_jsons = ground_truth["gt_parses"] |
|
else: |
|
assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict) |
|
gt_jsons = [ground_truth["gt_parse"]] |
|
|
|
self.gt_token_sequences.append( |
|
[ |
|
task_start_token |
|
+ self.donut_model.json2token( |
|
gt_json, |
|
update_special_tokens_for_json_key=self.split == "train", |
|
sort_json_key=self.sort_json_key, |
|
) |
|
+ self.donut_model.decoder.tokenizer.eos_token |
|
for gt_json in gt_jsons |
|
] |
|
) |
|
|
|
self.donut_model.decoder.add_special_tokens([self.task_start_token, self.prompt_end_token]) |
|
self.prompt_end_token_id = self.donut_model.decoder.tokenizer.convert_tokens_to_ids(self.prompt_end_token) |
|
|
|
def __len__(self) -> int: |
|
return self.dataset_length |
|
|
|
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
|
""" |
|
Load image from image_path of given dataset_path and convert into input_tensor and labels. |
|
Convert gt data into input_ids (tokenized string) |
|
|
|
Returns: |
|
input_tensor : preprocessed image |
|
input_ids : tokenized gt_data |
|
labels : masked labels (model doesn't need to predict prompt and pad token) |
|
""" |
|
sample = self.dataset[idx] |
|
|
|
|
|
input_tensor = self.donut_model.encoder.prepare_input(sample["image"], random_padding=self.split == "train") |
|
|
|
|
|
processed_parse = random.choice(self.gt_token_sequences[idx]) |
|
input_ids = self.donut_model.decoder.tokenizer( |
|
processed_parse, |
|
add_special_tokens=False, |
|
max_length=self.max_length, |
|
padding="max_length", |
|
truncation=True, |
|
return_tensors="pt", |
|
)["input_ids"].squeeze(0) |
|
|
|
if self.split == "train": |
|
labels = input_ids.clone() |
|
labels[ |
|
labels == self.donut_model.decoder.tokenizer.pad_token_id |
|
] = self.ignore_id |
|
labels[ |
|
: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1 |
|
] = self.ignore_id |
|
return input_tensor, input_ids, labels |
|
else: |
|
prompt_end_index = torch.nonzero( |
|
input_ids == self.prompt_end_token_id |
|
).sum() |
|
return input_tensor, input_ids, prompt_end_index, processed_parse |
|
|
|
|
|
class JSONParseEvaluator: |
|
""" |
|
Calculate n-TED(Normalized Tree Edit Distance) based accuracy and F1 accuracy score |
|
""" |
|
|
|
@staticmethod |
|
def flatten(data: dict): |
|
""" |
|
Convert Dictionary into Non-nested Dictionary |
|
Example: |
|
input(dict) |
|
{ |
|
"menu": [ |
|
{"name" : ["cake"], "count" : ["2"]}, |
|
{"name" : ["juice"], "count" : ["1"]}, |
|
] |
|
} |
|
output(list) |
|
[ |
|
("menu.name", "cake"), |
|
("menu.count", "2"), |
|
("menu.name", "juice"), |
|
("menu.count", "1"), |
|
] |
|
""" |
|
flatten_data = list() |
|
|
|
def _flatten(value, key=""): |
|
if type(value) is dict: |
|
for child_key, child_value in value.items(): |
|
_flatten(child_value, f"{key}.{child_key}" if key else child_key) |
|
elif type(value) is list: |
|
for value_item in value: |
|
_flatten(value_item, key) |
|
else: |
|
flatten_data.append((key, value)) |
|
|
|
_flatten(data) |
|
return flatten_data |
|
|
|
@staticmethod |
|
def update_cost(node1: Node, node2: Node): |
|
""" |
|
Update cost for tree edit distance. |
|
If both are leaf node, calculate string edit distance between two labels (special token '<leaf>' will be ignored). |
|
If one of them is leaf node, cost is length of string in leaf node + 1. |
|
If neither are leaf node, cost is 0 if label1 is same with label2 othewise 1 |
|
""" |
|
label1 = node1.label |
|
label2 = node2.label |
|
label1_leaf = "<leaf>" in label1 |
|
label2_leaf = "<leaf>" in label2 |
|
if label1_leaf == True and label2_leaf == True: |
|
return edit_distance(label1.replace("<leaf>", ""), label2.replace("<leaf>", "")) |
|
elif label1_leaf == False and label2_leaf == True: |
|
return 1 + len(label2.replace("<leaf>", "")) |
|
elif label1_leaf == True and label2_leaf == False: |
|
return 1 + len(label1.replace("<leaf>", "")) |
|
else: |
|
return int(label1 != label2) |
|
|
|
@staticmethod |
|
def insert_and_remove_cost(node: Node): |
|
""" |
|
Insert and remove cost for tree edit distance. |
|
If leaf node, cost is length of label name. |
|
Otherwise, 1 |
|
""" |
|
label = node.label |
|
if "<leaf>" in label: |
|
return len(label.replace("<leaf>", "")) |
|
else: |
|
return 1 |
|
|
|
def normalize_dict(self, data: Union[Dict, List, Any]): |
|
""" |
|
Sort by value, while iterate over element if data is list |
|
""" |
|
if not data: |
|
return {} |
|
|
|
if isinstance(data, dict): |
|
new_data = dict() |
|
for key in sorted(data.keys(), key=lambda k: (len(k), k)): |
|
value = self.normalize_dict(data[key]) |
|
if value: |
|
if not isinstance(value, list): |
|
value = [value] |
|
new_data[key] = value |
|
|
|
elif isinstance(data, list): |
|
if all(isinstance(item, dict) for item in data): |
|
new_data = [] |
|
for item in data: |
|
item = self.normalize_dict(item) |
|
if item: |
|
new_data.append(item) |
|
else: |
|
new_data = [str(item).strip() for item in data if type(item) in {str, int, float} and str(item).strip()] |
|
else: |
|
new_data = [str(data).strip()] |
|
|
|
return new_data |
|
|
|
def cal_f1(self, preds: List[dict], answers: List[dict]): |
|
""" |
|
Calculate global F1 accuracy score (field-level, micro-averaged) by counting all true positives, false negatives and false positives |
|
""" |
|
total_tp, total_fn_or_fp = 0, 0 |
|
for pred, answer in zip(preds, answers): |
|
pred, answer = self.flatten(self.normalize_dict(pred)), self.flatten(self.normalize_dict(answer)) |
|
for field in pred: |
|
if field in answer: |
|
total_tp += 1 |
|
answer.remove(field) |
|
else: |
|
total_fn_or_fp += 1 |
|
total_fn_or_fp += len(answer) |
|
return total_tp / (total_tp + total_fn_or_fp / 2) |
|
|
|
def construct_tree_from_dict(self, data: Union[Dict, List], node_name: str = None): |
|
""" |
|
Convert Dictionary into Tree |
|
|
|
Example: |
|
input(dict) |
|
|
|
{ |
|
"menu": [ |
|
{"name" : ["cake"], "count" : ["2"]}, |
|
{"name" : ["juice"], "count" : ["1"]}, |
|
] |
|
} |
|
|
|
output(tree) |
|
<root> |
|
| |
|
menu |
|
/ \ |
|
<subtree> <subtree> |
|
/ | | \ |
|
name count name count |
|
/ | | \ |
|
<leaf>cake <leaf>2 <leaf>juice <leaf>1 |
|
""" |
|
if node_name is None: |
|
node_name = "<root>" |
|
|
|
node = Node(node_name) |
|
|
|
if isinstance(data, dict): |
|
for key, value in data.items(): |
|
kid_node = self.construct_tree_from_dict(value, key) |
|
node.addkid(kid_node) |
|
elif isinstance(data, list): |
|
if all(isinstance(item, dict) for item in data): |
|
for item in data: |
|
kid_node = self.construct_tree_from_dict( |
|
item, |
|
"<subtree>", |
|
) |
|
node.addkid(kid_node) |
|
else: |
|
for item in data: |
|
node.addkid(Node(f"<leaf>{item}")) |
|
else: |
|
raise Exception(data, node_name) |
|
return node |
|
|
|
def cal_acc(self, pred: dict, answer: dict): |
|
""" |
|
Calculate normalized tree edit distance(nTED) based accuracy. |
|
1) Construct tree from dict, |
|
2) Get tree distance with insert/remove/update cost, |
|
3) Divide distance with GT tree size (i.e., nTED), |
|
4) Calculate nTED based accuracy. (= max(1 - nTED, 0 ). |
|
""" |
|
pred = self.construct_tree_from_dict(self.normalize_dict(pred)) |
|
answer = self.construct_tree_from_dict(self.normalize_dict(answer)) |
|
return max( |
|
0, |
|
1 |
|
- ( |
|
zss.distance( |
|
pred, |
|
answer, |
|
get_children=zss.Node.get_children, |
|
insert_cost=self.insert_and_remove_cost, |
|
remove_cost=self.insert_and_remove_cost, |
|
update_cost=self.update_cost, |
|
return_operations=False, |
|
) |
|
/ zss.distance( |
|
self.construct_tree_from_dict(self.normalize_dict({})), |
|
answer, |
|
get_children=zss.Node.get_children, |
|
insert_cost=self.insert_and_remove_cost, |
|
remove_cost=self.insert_and_remove_cost, |
|
update_cost=self.update_cost, |
|
return_operations=False, |
|
) |
|
), |
|
) |
|
|
