huggingface / huggingface/pixparse

unified abstractions for collate functions

Open
#26 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
25
Forks
4
PR merge metrics
No merged PRs in 30d

Description

In PR #24 , collate_fn is
* not defined in `task_cruller_pretrain.py`
* in `task_finetune_rvlcdip.py`, defined as
```python
def collate_fn(self, batch):
return CollateRVLCDIP(
self.tokenizer,
self.image_preprocess_train,
self.task_start_token,
self.label_int2str,
)
```
where
```python
class CollateRVLCDIP:
"""
basic collator for PIL images, as returned by rvlcdip dataloader (among others)
"""

def __init__(
self,
tokenizer,
image_preprocess,
start_token: str,
label_int2str: dict,
):
self.tokenizer = tokenizer
self.tokenizer_fn = lambda x: self.tokenizer(
x,
add_special_tokens=False,
return_tensors='pt',
max_length=5,
padding='max_length',
truncation=True).input_ids[0]
self.image_preprocess = image_preprocess
self.start_token = start_token
self.int2str = label_int2str

def __call__(self, batch):
images = [item["image"] for item in batch]
labels = [item["label"] for item in batch]
labels_tokens = [
self.tokenizer_fn(self.start_token + "<" + self.int2str[label] + "/>" + self.tokenizer.eos_token)
for label in labels
]
images = torch.stack([self.image_preprocess(img) for img in images])
labels = torch.stack(labels_tokens)
targets = torch.stack([text_input_to_target(text, self.tokenizer) for text in labels])
labels = labels[:, :-1]
targets = targets[:, 1:]

return {"image": images, "label": labels, "text_target": targets}
```

* In `task_finetune_cord.py`, defined directly as
```python
def collate_fn(self, batch):
"""
basic collator for PIL images, as returned by rvlcdip dataloader (among others)
"""
tokenizer_fn = lambda x: self.tokenizer(
x, # FIXME move this batcher/tokenizer elsewhere
add_special_tokens=False,
return_tensors="pt",
max_length=512,
padding="max_length",
truncation=True,
).input_ids[0]

images = [item["image"] for item in batch]
raw_texts = [literal_eval(item["ground_truth"])["gt_parse"] for item in batch]
inputs_to_stack = []
for text in raw_texts:
tokens_from_json, _ = json2token(text, self.tokenizer.all_special_tokens, sort_json_key=False)
inputs_to_stack.append(tokenizer_fn(
self.task_start_token
#+ self.tokenizer.bos_token
+ tokens_from_json
+ self.tokenizer.eos_token
))
text_inputs = torch.stack(
inputs_to_stack
)
targets = torch.stack([self.text_input_to_target(text) for text in text_inputs])

transform = self.image_preprocess_train
images = torch.stack([transform(img) for img in images])
text_inputs = text_inputs[:, :-1]
targets = targets[:, 1:]
return {
"image": images,
"label": text_inputs,
"text_target": targets,
}
```

And same for `task_finetune_docvqa` and the 3 eval tasks associated.

Suggestion:
Make a template Collator class with inherited cases per class, instantiate collate_fn from these

Like

```python
class CollateTask:
def __init__(self, tokenizer, image_preprocess, start_token):
self.tokenizer = tokenizer
self.image_preprocess = image_preprocess
self.start_token = start_token

def tokenizer_fn(self, x, max_length=512):
return self.tokenizer(
x,
add_special_tokens=False,
return_tensors="pt",
max_length=max_length,
padding="max_length",
truncation=True,
).input_ids[0]

def __call__(self, batch):
raise NotImplementedError("This method should be overridden by child classes")
```
And for RVLCDIP

```python
class CollateRVLCDIP(CollateTask):
def __init__(self, tokenizer, image_preprocess, start_token, label_int2str):
super().__init__(tokenizer, image_preprocess, start_token)
self.label_int2str = label_int2str

def __call__(self, batch):
....
```

And so on for CORD and DocVQA and others so that the collate_fn simply returns the instantiation of the collator class

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.