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="", prompt_end_token="", 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"" close_tag = f"" 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"".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 }