Views
No views yet
1import json
2from glob import glob
3from transformers import AutoProcessor, AutoModelForDocumentQuestionAnswering
4
5import torch
6import numpy as np
7
8model_name = "TusharGoel/LayoutLMv2-finetuned-docvqa"
9processor = AutoProcessor.from_pretrained(model_name)
10model = AutoModelForDocumentQuestionAnswering.from_pretrained(model_name)
11
12
13def pipeline(question, words, boxes, **kwargs):
14
15 images = kwargs["images"]
16 try:
17 encoding = processor(
18 images, question, words,boxes = boxes, return_token_type_ids=True, return_tensors="pt", truncation = True
19 )
20 word_ids = encoding.word_ids(0)
21
22 outputs = model(**encoding)
23
24 start_scores = outputs.start_logits
25 end_scores = outputs.end_logits
26
27
28 start, end = word_ids[start_scores.argmax(-1)], word_ids[end_scores.argmax(-1)]
29 answer = " ".join(words[start : end + 1])
30
31 start_scores, end_scores = start_scores.detach().numpy(), end_scores.detach().numpy()
32 undesired_tokens = encoding['attention_mask']
33 undesired_tokens_mask = undesired_tokens == 0.0
34
35 start_ = np.where(undesired_tokens_mask, -10000.0, start_scores)
36 end_ = np.where(undesired_tokens_mask, -10000.0, end_scores)
37 start_ = np.exp(start_ - np.log(np.sum(np.exp(start_), axis=-1, keepdims=True)))
38 end_ = np.exp(end_ - np.log(np.sum(np.exp(end_), axis=-1, keepdims=True)))
39
40 outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))
41 max_answer_len = 20
42 candidates = np.tril(np.triu(outer), max_answer_len - 1)
43 scores_flat = candidates.flatten()
44
45 idx_sort = [np.argmax(scores_flat)]
46 start, end = np.unravel_index(idx_sort, candidates.shape)[1:]
47
48 scores = candidates[0, start, end]
49 score = scores[0]
50 except Exception as e:
51 answer, score = "", 0.0
52 return answer, score