48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
from typing import Any
|
|
|
|
import torch
|
|
from transformers import DataCollatorForLanguageModeling
|
|
|
|
from texteller.constants import MAX_TOKEN_SIZE, MIN_HEIGHT, MIN_WIDTH
|
|
|
|
|
|
def _left_move(x: torch.Tensor, pad_val):
|
|
assert len(x.shape) == 2, "x should be 2-dimensional"
|
|
lefted_x = torch.ones_like(x)
|
|
lefted_x[:, :-1] = x[:, 1:]
|
|
lefted_x[:, -1] = pad_val
|
|
return lefted_x
|
|
|
|
|
|
def tokenize_fn(samples: dict[str, list[Any]], tokenizer=None) -> dict[str, list[Any]]:
|
|
assert tokenizer is not None, "tokenizer should not be None"
|
|
tokenized_formula = tokenizer(samples["latex_formula"], return_special_tokens_mask=True)
|
|
tokenized_formula["pixel_values"] = samples["image"]
|
|
return tokenized_formula
|
|
|
|
|
|
def collate_fn(samples: list[dict[str, Any]], tokenizer=None) -> dict[str, list[Any]]:
|
|
assert tokenizer is not None, "tokenizer should not be None"
|
|
pixel_values = [dic.pop("pixel_values") for dic in samples]
|
|
|
|
clm_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
|
|
|
|
batch = clm_collator(samples)
|
|
batch["pixel_values"] = pixel_values
|
|
batch["decoder_input_ids"] = batch.pop("input_ids")
|
|
batch["decoder_attention_mask"] = batch.pop("attention_mask")
|
|
|
|
batch["labels"] = _left_move(batch["labels"], -100)
|
|
|
|
# convert list of Image to a tensor with (B, C, H, W)
|
|
batch["pixel_values"] = torch.stack(batch["pixel_values"], dim=0)
|
|
return batch
|
|
|
|
|
|
def filter_fn(sample, tokenizer=None) -> bool:
|
|
return (
|
|
sample["image"].height > MIN_HEIGHT
|
|
and sample["image"].width > MIN_WIDTH
|
|
and len(tokenizer(sample["latex_formula"])["input_ids"]) < MAX_TOKEN_SIZE - 10
|
|
)
|