Spaces:
Running
Running
File size: 8,984 Bytes
8af6af2 |
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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 |
import json
from typing import Any, Dict, Optional
import numpy as np
from PIL import Image
from datasets import DatasetDict
from torch.utils.data import Dataset
from transformers import pipeline, DonutProcessor
class DonutFinetuned:
DEFAULT_PIPELINE = pipeline(
task="image-to-text",
model="naver-clova-ix/donut-base"
)
@classmethod
def predict(cls, image: np.ndarray) -> dict:
image = Image.fromarray(image)
result = cls.DEFAULT_PIPELINE(image)
return result
class DonutDatasets:
"""
Modified from:
https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Donut/CORD/Fine_tune_Donut_on_a_custom_dataset_(CORD)_with_PyTorch_Lightning.ipynb
Donut PyTorch Dataset Wrapper (supports train/validation/test splits)
- Dynamic field names and JSON-to-token conversion
- Returns PyTorch Datasets with __getitem__ producing tensors
- Splits controlled by train_split/validation_split/test_split
- Only single JSON annotation supported
- Supports subscripting: datasets["train"], datasets["validation"], datasets["test"]
Args:
- datasets: DatasetDict containing train/validation/test splits
- processor: DonutProcessor for image processing
- image_column: Column name for images in the dataset
- annotation_column: Column name for annotations in the dataset
- task_start_token: Token to start the task
- prompt_end_token: Token to end the prompt
- max_length: Maximum length of tokenized sequences
- train_split: Fraction of data to use for training (0.0-1.0)
- validation_split: Fraction of data to use for validation (0.0-1.0)
- test_split: Fraction of data to use for testing (0.0-1.0)
- ignore_index: Index to ignore in labels (default: -100)
- sort_json_key: Whether to sort JSON keys (default: True)
- seed: Random seed for reproducibility. If None, use OS random seed (default: None)
- shuffle: Whether to shuffle the dataset (default: True)
Returns:
- DonutDatasets object with train/validation/test splits
Example:
datasets = DonutDatasets(
datasets=dataset_dict,
processor=processor,
image_column="image",
annotation_column="annotation",
task_start_token="<s_task>",
prompt_end_token="<s_prompt>",
max_length=512,
train_split=0.8,
validation_split=0.1,
test_split=0.1
)
train_dataset = datasets["train"]
validation_dataset = datasets["validation"]
test_dataset = datasets["test"]
Note:
- The dataset must be a DatasetDict with train/validation/test splits
- The processor must be a DonutProcessor instance
- The image_column and annotation_column must exist in the dataset
- The task_start_token and prompt_end_token must be unique tokens
- The max_length should be set according to the model's maximum input length
- The ignore_index is used for padding in labels (default: -100)
- The sort_json_key option determines whether JSON keys are sorted or not
"""
def __init__(
self,
datasets: DatasetDict,
processor: DonutProcessor,
image_column: str,
annotation_column: str,
task_start_token: str,
prompt_end_token: str,
max_length: int = 512,
train_split: float = 1.0,
validation_split: float = 0.0,
test_split: float = 0.0,
ignore_index: int = -100,
sort_json_key: bool = True,
seed: Optional[int] = None,
shuffle: bool = True
):
assert abs(train_split + validation_split + test_split - 1.0) < 1e-6, (
"train/validation/test splits must sum to 1"
)
self.processor = processor
self.tokenizer = processor.tokenizer
self.image_column = image_column
self.annotation_column = annotation_column
self.max_length = max_length
self.task_start_token = task_start_token
self.prompt_end_token = prompt_end_token or task_start_token
self.ignore_index = ignore_index
self.sort_json_key = sort_json_key
# Perform split on provided datasets
raw = datasets
parts: Dict[str, Any] = {}
if train_split < 1.0:
split1 = raw["train"].train_test_split(test_size=1 - train_split, seed=seed, shuffle=shuffle)
parts["train"] = split1["train"]
rest = split1["test"]
if validation_split > 0:
val_frac = validation_split / (validation_split + test_split)
split2 = rest.train_test_split(test_size=1 - val_frac, seed=seed, shuffle=shuffle)
parts["validation"] = split2["train"]
parts["test"] = split2["test"]
else:
parts["test"] = rest
else:
parts = dict(raw)
# Create individual split datasets
self._splits: Dict[str, Dataset] = {}
for name, ds in parts.items():
self._splits[name] = _SplitDataset(
hf_dataset=ds,
processor=self.processor,
image_column=self.image_column,
annotation_column=self.annotation_column,
max_length=self.max_length,
ignore_index=self.ignore_index,
sort_json_key=self.sort_json_key,
task_start_token=self.task_start_token,
prompt_end_token=self.prompt_end_token,
)
def __getitem__(self, split: str) -> Dataset:
"""
Return the dataset split by name, e.g., datasets["train"]
"""
if split in self._splits:
return self._splits[split]
raise KeyError(f"Unknown split '{split}'. Available splits: {list(self._splits.keys())}")
def __repr__(self):
return f"DonutDatasets(splits={list(self._splits.keys())})"
class _SplitDataset(Dataset):
"""
PyTorch Dataset for a single split, returns (pixel_values, labels, target_sequence)
"""
def __init__(
self,
hf_dataset,
processor: DonutProcessor,
image_column: str,
annotation_column: str,
max_length: int,
ignore_index: int,
sort_json_key: bool,
task_start_token: str,
prompt_end_token: str,
):
self.processor = processor
self.tokenizer = processor.tokenizer
self.hf_dataset = hf_dataset
self.image_column = image_column
self.annotation_column = annotation_column
self.max_length = max_length
self.ignore_index = ignore_index
self.sort_json_key = sort_json_key
self.task_start_token = task_start_token
self.prompt_end_token = prompt_end_token
# Prepare tokenized ground-truth sequences (single annotation)
self.gt_token_sequences = []
for sample in self.hf_dataset:
gt = sample[self.annotation_column]
if isinstance(gt, str):
gt = json.loads(gt)
seq = self._json_to_token(gt) + self.tokenizer.eos_token
self.gt_token_sequences.append(seq)
# Add special tokens to tokenizer
self.tokenizer.add_tokens([self.task_start_token, self.prompt_end_token])
def _json_to_token(self, obj: Any) -> str:
if isinstance(obj, dict):
keys = sorted(obj.keys()) if self.sort_json_key else obj.keys()
seq = ""
for k in keys:
open_tag = f"<s_{k}>"
close_tag = f"</s_{k}>"
self.tokenizer.add_special_tokens({"additional_special_tokens": [open_tag, close_tag]})
seq += open_tag + self._json_to_token(obj[k]) + close_tag
return seq
if isinstance(obj, list):
return r"<sep/>".join(self._json_to_token(x) for x in obj)
return str(obj)
def __len__(self):
return len(self.hf_dataset)
def __getitem__(self, idx: int):
sample = self.hf_dataset[idx]
pixel_values = self.processor(sample[self.image_column], return_tensors="pt").pixel_values.squeeze()
target_seq = self.gt_token_sequences[idx]
tokens = self.tokenizer(
target_seq,
add_special_tokens=False,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
input_ids = tokens.input_ids.squeeze(0)
labels = input_ids.clone()
labels[labels == self.tokenizer.pad_token_id] = self.ignore_index
return {
"pixel_values": pixel_values,
"input_ids": input_ids,
"attention_mask": tokens.attention_mask.squeeze(0),
"labels": labels,
"target_sequence": target_seq
} |