upd: generate_answer

app.py

@@ -5,8 +5,14 @@ from transformers import (
     AutoProcessor,
     AutoModelForDocumentQuestionAnswering,
 )
-from transformers import pipeline
+import torch
 
+processor = AutoProcessor.from_pretrained(
+    "MariaK/layoutlmv2-base-uncased_finetuned_docvqa_v2"
+)
+model = AutoModelForDocumentQuestionAnswering.from_pretrained(
+    "MariaK/layoutlmv2-base-uncased_finetuned_docvqa_v2"
+)
 
 tokenizer_ru2en = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
 model_ru2en = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
@@ -14,14 +20,6 @@ model_ru2en = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ru-en"
 tokenizer_en2ru = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
 model_en2ru = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
 
-git_processor_base = AutoProcessor.from_pretrained("microsoft/layoutlmv2-base-uncased")
-
-image_processor = git_processor_base.image_processor
-
-
-def preprocess_image(image):
-    return git_processor_base(images=image, return_tensors="pt").pixel_values
-
 
 def translate_ru2en(text):
     inputs = tokenizer_ru2en(text, return_tensors="pt")
@@ -39,11 +37,23 @@ def translate_en2ru(text):
 
 def generate_answer_git(image, question):
 
-    qa_pipeline = pipeline(
-        "document-question-answering",
-        model="andgrt/layoutlmv2-base-uncased_finetuned_docvqa",
+    with torch.no_grad():
+        encoding = processor(
+            images=image,
+            text=question,
+            return_tensors="pt",
+            max_length=512,
+            truncation=True,
+        )
+        outputs = model(**encoding)
+        start_logits = outputs.start_logits
+        end_logits = outputs.end_logits
+        predicted_start_idx = start_logits.argmax(-1).item()
+        predicted_end_idx = end_logits.argmax(-1).item()
+
+    return processor.tokenizer.decode(
+        encoding.input_ids.squeeze()[predicted_start_idx : predicted_end_idx + 1]
     )
-    return qa_pipeline(preprocess_image(image), question)[0]["answer"]
 
 
 def generate_answer(image, question):