app.py

import gradio as gr
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoProcessor,
    AutoModelForDocumentQuestionAnswering,
)
from transformers import pipeline
import torch

tokenizer_ru2en = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
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")

git_model_base = AutoModelForDocumentQuestionAnswering.from_pretrained(
    "andgrt/layoutlmv2-base-uncased_finetuned_docvqa"
)

device = "cuda" if torch.cuda.is_available() else "cpu"
git_model_base.to(device)
image_processor = git_processor_base.image_processor


def preprocess_image(image):
    """Преобразуем изображение для модели"""
    image_rgb = image.convert("RGB")
    return image_processor(image_rgb, return_tensors="pt").pixel_values.to(device)


def translate_ru2en(text):
    inputs = tokenizer_ru2en(text, return_tensors="pt")
    outputs = model_ru2en.generate(**inputs)
    translated_text = tokenizer_ru2en.decode(outputs[0], skip_special_tokens=True)
    return translated_text


def translate_en2ru(text):
    inputs = tokenizer_en2ru(text, return_tensors="pt")
    outputs = model_en2ru.generate(**inputs)
    translated_text = tokenizer_en2ru.decode(outputs[0], skip_special_tokens=True)
    return translated_text


def generate_answer_git(image, question):

    qa_pipeline = pipeline(
        "document-question-answering",
        model="andgrt/layoutlmv2-base-uncased_finetuned_docvqa",
    )
    return qa_pipeline(preprocess_image(image), question)[0]["answer"]
    # pixel_values, _, _ = preprocess_image(image)
    # input_ids = processor(text=question, add_special_tokens=False).input_ids
    # input_ids = [processor.tokenizer.cls_token_id] + input_ids
    # input_ids = torch.tensor(input_ids).unsqueeze(0).to(device)

    # generated_ids = model.generate(
    #     pixel_values=pixel_values, input_ids=input_ids, max_length=50
    # )
    # generated_answer = processor.batch_decode(generated_ids, skip_special_tokens=True)
    # return generated_answer[0]


def generate_answer(image, question):
    question_en = translate_ru2en(question)
    print(f"Вопрос на английском: {question_en}")

    answer_en = generate_answer_git(image, question_en)
    print(f"Ответ на английском: {answer_en}")

    answer_ru = translate_en2ru(answer_en)

    return answer_ru


examples = [
    ["doc.png", "О чем данный документ?"],
]

interface = gr.Interface(
    fn=generate_answer,
    inputs=[
        gr.Image(type="pil"),
        gr.Textbox(label="Вопрос (на русском)", placeholder="Ваш вопрос"),
    ],
    outputs=gr.Textbox(label="Ответ (на русском)"),
    examples=examples,
    title="Демо визуального ответчика на вопросы (на русском)",
    description=(
        "Gradio демо для модели doc-qa с переводом вопросов и ответов"
        "на русский язык. Загрузите изображение и задайте вопрос, чтобы"
        "получить ответ. Вы также можете использовать голосовой ввод!"
    ),
    allow_flagging="never",
)

interface.launch(debug=True, share=True)