init

app.py

@@ -0,0 +1,97 @@
+import gradio as gr
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    AutoProcessor,
+    AutoModelForCausalLM,
+)
+import torch
+import pyttsx3
+
+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(
+    "andgrt/layoutlmv2-base-uncased_finetuned_docvqa"
+)
+git_model_base = AutoModelForCausalLM.from_pretrained(
+    "andgrt/layoutlmv2-base-uncased_finetuned_docvqa"
+)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+git_model_base.to(device)
+
+
+engine = pyttsx3.init()
+
+
+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(processor, model, image, question):
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
+    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(
+        git_processor_base, git_model_base, image, question_en
+    )
+    print(f"Ответ на английском: {answer_en}")
+
+    answer_ru = translate_en2ru(answer_en)
+
+    engine.say(answer_ru)
+    engine.runAndWait()
+
+    return answer_ru
+
+
+examples = [
+    ["doc.png", "О чем данный документ?"],
+]
+
+interface = gr.Interface(
+    fn=generate_answer,
+    inputs=[
+        gr.inputs.Image(type="pil"),
+        gr.inputs.Textbox(label="Вопрос (на русском)", placeholder="Ваш вопрос"),
+    ],
+    outputs=gr.outputs.Textbox(label="Ответ (на русском)"),
+    examples=examples,
+    title="Демо визуального ответчика на вопросы (на русском)",
+    description=(
+        "Gradio демо для модели doc-qa с переводом вопросов и ответов"
+        "на русский язык. Загрузите изображение и задайте вопрос, чтобы"
+        "получить ответ. Вы также можете использовать голосовой ввод!"
+    ),
+    allow_flagging="never",
+    enable_queue=True,
+)
+
+interface.launch(debug=True, share=True)