Merge pull request #1 from NicBonetto/initialize-model

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+core/__pycache__/
+trocr-ocr/

config.py

@@ -0,0 +1,8 @@
+DATASET_NAME = 'Teklia/IAM-line'
+MODEL_NAME = 'microsoft/trocr-small-printed'
+TRAIN_SPLIT = 'train[:200]'
+TEST_SPLIT_RATIO = 0.2
+BATCH_SIZE = 2
+EPOCHS = 1
+OUTPUT_DIR = './trocr-ocr'
+

core/data.py

@@ -0,0 +1,27 @@
+import torch
+from datasets import load_dataset
+from config import DATASET_NAME, TRAIN_SPLIT, TEST_SPLIT_RATIO
+from core.model import processor
+
+def preprocess_batch(batch):
+    images = [img.convert("RGB") for img in batch["image"]]
+    pixel_values = processor(images=images, return_tensors="pt").pixel_values
+    batch["pixel_values"] = pixel_values
+
+    labels = processor.tokenizer(batch["text"], padding=True, truncation=True).input_ids
+    batch["labels"] = labels
+
+    return batch
+
+def load():
+    dataset = load_dataset(DATASET_NAME, split = TRAIN_SPLIT)
+    train_test = dataset.train_test_split(test_size = TEST_SPLIT_RATIO)
+    train_ds = train_test['train']
+    eval_ds = train_test['test']
+
+    train_ds = train_ds.map(preprocess_batch, batched = True, remove_columns = train_ds.column_names)
+    eval_ds = eval_ds.map(preprocess_batch, batched = True, remove_columns = eval_ds.column_names)
+
+    return train_ds, eval_ds
+
+

core/model.py

@@ -0,0 +1,13 @@
+import torch
+from transformers import VisionEncoderDecoderModel, TrOCRProcessor
+from config import MODEL_NAME
+
+device = torch.device('mps') if torch.backends.mps.is_available() else torch.device('cpu')
+processor = TrOCRProcessor.from_pretrained(MODEL_NAME, use_fast = False)
+model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME).to(device)
+
+model.config.decoder_start_token_id = processor.tokenizer.bos_token_id
+
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+if model.config.pad_token_id is None:
+    model.config.pad_token_id = processor.tokenizer.eos_token_id

core/train.py

@@ -0,0 +1,32 @@
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
+from core.model import model, processor
+from core.data import load
+from core.utils import compute_metrics
+from config import OUTPUT_DIR, BATCH_SIZE, EPOCHS
+
+train_ds, eval_ds = load()
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir = OUTPUT_DIR,
+    per_device_train_batch_size = BATCH_SIZE,
+    per_device_eval_batch_size = BATCH_SIZE,
+    predict_with_generate = True,
+    eval_strategy = 'epoch',
+    logging_steps = 10,
+    num_train_epochs = EPOCHS,
+    save_total_limit = 1,
+    fp16 = False
+)
+
+trainer = Seq2SeqTrainer(
+    model = model,
+    args = training_args,
+    train_dataset = train_ds,
+    eval_dataset = eval_ds,
+    processing_class = processor.image_processor,
+    data_collator = default_data_collator,
+    compute_metrics = compute_metrics
+)
+
+if __name__ == '__main__':
+    trainer.train()

core/utils.py

@@ -0,0 +1,15 @@
+import evaluate
+from core.model import processor
+
+cer_metric = evaluate.load('cer')
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens = True)
+    label_ids[ label_ids == -100 ] = processor.tokenizer.pad_token_id
+    label_str = processor.batch_decode(label_ids, skip_special_tokens = True)
+    cer = cer_metric.compute(predictions = pred_str, references = label_str)
+
+    return { 'cer': cer }
+

requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchvision
+transformers
+datasets
+evaluate
+jiwer
+Pillow
+