<jupyter_start><jupyter_text>In this notebook, we are going to fine-tune `LayoutLMv2ForTokenClassification` on the [CORD](https://github.com/clovaai/cord) dataset. The goal for the model is to label words appearing in scanned documents (namely, receipts) appropriately. This task is treated as a NER problem (sequence labeling). However, compared to BERT, LayoutLMv2 also incorporates visual and layout information about the tokens when encoding them into vectors. This makes the LayoutLMv2 model very powerful for document understanding tasks.LayoutLMv2 is itself an upgrade of LayoutLM. The main novelty of LayoutLMv2 is that it also pre-trains visual embeddings, whereas the original LayoutLM only adds visual embeddings during fine-tuning.* Paper: https://arxiv.org/abs/2012.14740* Original repo: https://github.com/microsoft/unilm/tree/master/layoutlmv2NOTES: * you first need to prepare the CORD dataset for LayoutLMv2. For that, check out the notebook "Prepare CORD for LayoutLMv2".* this notebook is heavily inspired by [this Github repository](https://github.com/omarsou/layoutlm_CORD), which fine-tunes both BERT and LayoutLM (v1) on the CORD dataset. Install dependenciesFirst, we install the required libraries:* Transformers (for the LayoutLMv2 model)* Datasets (for data preprocessing)* Seqeval (for metrics)* Detectron2 (which LayoutLMv2 requires for its visual backbone).<jupyter_code>!pip install transformers
!pip install -q datasets seqeval
!pip install pyyaml==5.1
!pip install torch==1.10.0+cu111 torchvision==0.11.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

!python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.10/index.html
from google.colab import drive
drive.mount('/content/drive')<jupyter_output>Mounted at /content/drive<jupyter_text>Prepare the dataFirst, let's read in the annotations which we prepared in the other notebook. These contain the word-level annotations (words, labels, normalized bounding boxes).<jupyter_code>import pandas as pd

train = pd.read_pickle('/content/drive/MyDrive/sparrow/CORD/CORD_layoutlmv2_format/train.pkl')
val = pd.read_pickle('/content/drive/MyDrive/sparrow/CORD/CORD_layoutlmv2_format/dev.pkl')
test = pd.read_pickle('/content/drive/MyDrive/sparrow/CORD/CORD_layoutlmv2_format/test.pkl')<jupyter_output><empty_output><jupyter_text>Let's define a list of all unique labels. For that, let's first count the number of occurrences for each label:<jupyter_code>from collections import Counter

all_labels = [item for sublist in train[1] for item in sublist] + [item for sublist in val[1] for item in sublist] + [item for sublist in test[1] for item in sublist]
Counter(all_labels)<jupyter_output><empty_output><jupyter_text>As we can see, there are some labels that contain very few examples. Let's replace them by the "neutral" label "O" (which stands for "Outside").<jupyter_code>replacing_labels = {'menu.etc': 'O', 'mneu.itemsubtotal': 'O', 'menu.sub_etc': 'O', 'menu.sub_unitprice': 'O', 'menu.vatyn': 'O',
                  'void_menu.nm': 'O', 'void_menu.price': 'O', 'sub_total.othersvc_price': 'O'}
def replace_elem(elem):
  try:
    return replacing_labels[elem]
  except KeyError:
    return elem
def replace_list(ls):
  return [replace_elem(elem) for elem in ls]
train[1] = [replace_list(ls) for ls in train[1]]
val[1] = [replace_list(ls) for ls in val[1]]
test[1] = [replace_list(ls) for ls in test[1]]
all_labels = [item for sublist in train[1] for item in sublist] + [item for sublist in val[1] for item in sublist] + [item for sublist in test[1] for item in sublist]
Counter(all_labels)<jupyter_output><empty_output><jupyter_text>Now we have to save all the unique labels in a list.<jupyter_code>labels = list(set(all_labels))
print(labels)
import pickle

label2id = {label: idx for idx, label in enumerate(labels)}
id2label = {idx: label for idx, label in enumerate(labels)}
print(label2id)
print(id2label)

with open('/content/drive/MyDrive/sparrow/CORD/label2id.pkl', 'wb') as t:
    pickle.dump(label2id, t)

with open('/content/drive/MyDrive/sparrow/CORD/id2label.pkl', 'wb') as t:
    pickle.dump(id2label, t)
from os import listdir
from torch.utils.data import Dataset
import torch
from PIL import Image

class CORDDataset(Dataset):
    """CORD dataset."""

    def __init__(self, annotations, image_dir, processor=None, max_length=512):
        """
        Args:
            annotations (List[List]): List of lists containing the word-level annotations (words, labels, boxes).
            image_dir (string): Directory with all the document images.
            processor (LayoutLMv2Processor): Processor to prepare the text + image.
        """
        self.words, self.labels, self.boxes = annotations
        self.image_dir = image_dir
        self.image_file_names = [f for f in listdir(image_dir)]
        self.processor = processor

    def __len__(self):
        return len(self.image_file_names)

    def __getitem__(self, idx):
        # first, take an image
        item = self.image_file_names[idx]
        image = Image.open(self.image_dir + item).convert("RGB")

        # get word-level annotations 
        words = self.words[idx]
        boxes = self.boxes[idx]
        word_labels = self.labels[idx]

        assert len(words) == len(boxes) == len(word_labels)
        
        word_labels = [label2id[label] for label in word_labels]
        # use processor to prepare everything
        encoded_inputs = self.processor(image, words, boxes=boxes, word_labels=word_labels, 
                                        padding="max_length", truncation=True, 
                                        return_tensors="pt")
        
        # remove batch dimension
        for k,v in encoded_inputs.items():
          encoded_inputs[k] = v.squeeze()

        assert encoded_inputs.input_ids.shape == torch.Size([512])
        assert encoded_inputs.attention_mask.shape == torch.Size([512])
        assert encoded_inputs.token_type_ids.shape == torch.Size([512])
        assert encoded_inputs.bbox.shape == torch.Size([512, 4])
        assert encoded_inputs.image.shape == torch.Size([3, 224, 224])
        assert encoded_inputs.labels.shape == torch.Size([512]) 
      
        return encoded_inputs
from transformers import LayoutLMv2Processor

processor = LayoutLMv2Processor.from_pretrained("microsoft/layoutlmv2-base-uncased", revision="no_ocr")

train_dataset = CORDDataset(annotations=train,
                            image_dir='/content/drive/MyDrive/sparrow/CORD/train/image/', 
                            processor=processor)
val_dataset = CORDDataset(annotations=val,
                            image_dir='/content/drive/MyDrive/sparrow/CORD/dev/image/', 
                            processor=processor)
test_dataset = CORDDataset(annotations=test,
                            image_dir='/content/drive/MyDrive/sparrow/CORD/test/image/', 
                            processor=processor)<jupyter_output><empty_output><jupyter_text>Let's verify an example:<jupyter_code>encoding = train_dataset[0]
encoding.keys()
for k,v in encoding.items():
  print(k, v.shape)
print(processor.tokenizer.decode(encoding['input_ids']))
train[0][0]
train[1][0]
[id2label[label] for label in encoding['labels'].tolist() if label != -100]
for id, label in zip(encoding['input_ids'][:30], encoding['labels'][:30]):
  print(processor.tokenizer.decode([id]), label.item())<jupyter_output>[CLS] -100
2 10
apple 21
pie 21
22 22
, -100
000 -100
1 10
apple 21
cream 21
##chee -100
##se -100
pastry 21
13 22
, -100
000 -100
total 19
35 19
, -100
000 -100
cash 12
100 12
, -100
000 -100
change 17
65 17
, -100
000 -100
[SEP] -100
[PAD] -100<jupyter_text>Next, we create corresponding dataloaders.<jupyter_code>from torch.utils.data import DataLoader

train_dataloader = DataLoader(train_dataset, batch_size=2, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=2, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=2)<jupyter_output><empty_output><jupyter_text>Train the modelLet's train the model using native PyTorch. We use the AdamW optimizer with learning rate = 5e-5 (this is a good default value when fine-tuning Transformer-based models).<jupyter_code>from transformers import LayoutLMv2ForTokenClassification, AdamW
import torch
from tqdm.notebook import tqdm

model = LayoutLMv2ForTokenClassification.from_pretrained('microsoft/layoutlmv2-base-uncased',
                                                                      num_labels=len(labels))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
optimizer = AdamW(model.parameters(), lr=5e-5)

global_step = 0
num_train_epochs = 4

#put the model in training mode
model.train() 
for epoch in range(num_train_epochs):  
   print("Epoch:", epoch)
   for batch in tqdm(train_dataloader):
        # get the inputs;
        input_ids = batch['input_ids'].to(device)
        bbox = batch['bbox'].to(device)
        image = batch['image'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        token_type_ids = batch['token_type_ids'].to(device)
        labels = batch['labels'].to(device)

        # zero the parameter gradients
        optimizer.zero_grad()
        
        # forward + backward + optimize
        outputs = model(input_ids=input_ids,
                        bbox=bbox,
                        image=image,
                        attention_mask=attention_mask,
                        token_type_ids=token_type_ids,
                        labels=labels) 
        loss = outputs.loss
        
        # print loss every 100 steps
        if global_step % 100 == 0:
          print(f"Loss after {global_step} steps: {loss.item()}")

        loss.backward()
        optimizer.step()
        global_step += 1

model.save_pretrained("/content/drive/MyDrive/sparrow/CORD/Checkpoints")<jupyter_output><empty_output><jupyter_text>EvaluationLet's evaluate the model on the test set. First, let's do a sanity check on the first example of the test set.<jupyter_code># encoding = test_dataset[0]
# processor.tokenizer.decode(encoding['input_ids'])
# ground_truth_labels = [id2label[label] for label in encoding['labels'].squeeze().tolist() if label != -100]
# print(ground_truth_labels)
# for k,v in encoding.items():
#   encoding[k] = v.unsqueeze(0).to(device)

# model.eval()
# # forward pass
# outputs = model(input_ids=encoding['input_ids'], attention_mask=encoding['attention_mask'],
#                 token_type_ids=encoding['token_type_ids'], bbox=encoding['bbox'],
#                 image=encoding['image'])
# prediction_indices = outputs.logits.argmax(-1).squeeze().tolist()
# print(prediction_indices)
# prediction_indices = outputs.logits.argmax(-1).squeeze().tolist()
# predictions = [id2label[label] for gt, label in zip(encoding['labels'].squeeze().tolist(), prediction_indices) if gt != -100]
# print(predictions)
import numpy as np

preds_val = None
out_label_ids = None

# put model in evaluation mode
model.eval()
for batch in tqdm(test_dataloader, desc="Evaluating"):
    with torch.no_grad():
        input_ids = batch['input_ids'].to(device)
        bbox = batch['bbox'].to(device)
        image = batch['image'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        token_type_ids = batch['token_type_ids'].to(device)
        labels = batch['labels'].to(device)

        # forward pass
        outputs = model(input_ids=input_ids, bbox=bbox, image=image, attention_mask=attention_mask, 
                        token_type_ids=token_type_ids, labels=labels)
        
        if preds_val is None:
          preds_val = outputs.logits.detach().cpu().numpy()
          out_label_ids = batch["labels"].detach().cpu().numpy()
        else:
          preds_val = np.append(preds_val, outputs.logits.detach().cpu().numpy(), axis=0)
          out_label_ids = np.append(
              out_label_ids, batch["labels"].detach().cpu().numpy(), axis=0
          )
import warnings
warnings.filterwarnings("ignore")
from seqeval.metrics import (
    classification_report,
    f1_score,
    precision_score,
    recall_score)

def results_test(preds, out_label_ids, labels):
  preds = np.argmax(preds, axis=2)

  label_map = {i: label for i, label in enumerate(labels)}

  out_label_list = [[] for _ in range(out_label_ids.shape[0])]
  preds_list = [[] for _ in range(out_label_ids.shape[0])]

  for i in range(out_label_ids.shape[0]):
      for j in range(out_label_ids.shape[1]):
          if out_label_ids[i, j] != -100:
              out_label_list[i].append(label_map[out_label_ids[i][j]])
              preds_list[i].append(label_map[preds[i][j]])

  results = {
      "precision": precision_score(out_label_list, preds_list),
      "recall": recall_score(out_label_list, preds_list),
      "f1": f1_score(out_label_list, preds_list),
  }
  return results, classification_report(out_label_list, preds_list)
labels = list(set(all_labels))
val_result, class_report = results_test(preds_val, out_label_ids, labels)
print("Overall results:", val_result)
print(class_report)<jupyter_output>Overall results: {'precision': 0.9069069069069069, 'recall': 0.9158453373768006, 'f1': 0.9113542059600152}
                         precision    recall  f1-score   support

                enu.cnt       0.92      1.00      0.96       224
      enu.discountprice       0.20      0.20      0.20        10
       enu.itemsubtotal       0.00      0.00      0.00         6
                 enu.nm       0.92      0.95      0.93       251
                enu.num       0.89      0.73      0.80        11
              enu.price       0.97      0.97      0.97       247
            enu.sub_cnt       1.00      0.06      0.11        17
             enu.sub_nm       0.57      0.81      0.67        32
          enu.sub_price       0.73      0.80      0.76        20
          enu.unitprice       0.97      0.99      0.98        68
         otal.cashprice       0.95      0.82      0.88        71
       otal.changeprice       0.89      0.92      0.90        59
   otal.creditcardprice       0.78      0.82   [...]
{'fraction_non_alphanumeric': 0.11075725406935598, 'fraction_numerical': 0.03205944798301486, 'mean_word_length': 3.1684365781710913, 'pattern_counts': {'":': 3, '<': 34, '<?xml version=': 0, '>': 34, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '156743', 'license': 'Apache-2.0', 'n_tokens_mistral': 4723, 'n_tokens_neox': 4382, 'n_words': 1096, 'path': 'sparrow-research/app/Fine_tuning_LayoutLMv2ForTokenClassification_on_CORD.ipynb', 'repo_name': 'katanaml/sparrow', 'size': '14130'}