Skip to content

Instantly share code, notes, and snippets.

@microcoder-py
Last active March 7, 2022 11:38
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save microcoder-py/fd20745f5cc937bd483da2e588204b5d to your computer and use it in GitHub Desktop.
Save microcoder-py/fd20745f5cc937bd483da2e588204b5d to your computer and use it in GitHub Desktop.
from transformers import LayoutLMTokenizer, LayoutLMForSequenceClassification
import torch
tokenizer = LayoutLMTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
model = LayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased")
words = ["Hello", "world"]
normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
token_boxes = []
for word, box in zip(words, normalized_word_boxes):
word_tokens = tokenizer.tokenize(word)
token_boxes.extend([box] * len(word_tokens))
# add bounding boxes of cls + sep tokens
token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
encoding = tokenizer(" ".join(words), return_tensors="pt")
input_ids = encoding["input_ids"]
attention_mask = encoding["attention_mask"]
token_type_ids = encoding["token_type_ids"]
bbox = torch.tensor([token_boxes])
sequence_label = torch.tensor([1])
outputs = model(
input_ids=input_ids,
bbox=bbox,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
labels=sequence_label,
)
loss = outputs.loss
logits = outputs.logits
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment