DocQA-Demo / app.py
shurik-p's picture
Update app.py
1e8ff31 verified
try:
import detectron2
except:
import os
os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr
qa_pipeline = pipeline("document-question-answering", model="shurik-p/llmv2-docvqa-finetuned")
ru_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
ru_en_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
en_ru_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
def translate_ru2en(ru_question, ru_en_model, ru_tokenizer):
input_ids = ru_tokenizer.encode(ru_question, return_tensors="pt")
output_ids = ru_en_model.generate(input_ids, max_new_tokens=512)
en_question = ru_tokenizer.decode(output_ids[0], skip_special_tokens=True)
return en_question
def translate_en2ru(en_answer, en_ru_model, en_tokenizer):
input_ids = en_tokenizer.encode(en_answer, return_tensors="pt")
output_ids = en_ru_model.generate(input_ids, max_new_tokens=512)
ru_answer = en_tokenizer.decode(output_ids[0], skip_special_tokens=True)
return ru_answer
def ru_inference(image, ru_question):
en_question = translate_ru2en(ru_question, ru_en_model, ru_tokenizer)
en_answer = qa_pipeline(image=image, question=en_question)[0]['answer']
ru_answer = translate_en2ru(en_answer, en_ru_model, en_tokenizer)
return ru_answer
interface = gr.Interface(
fn=ru_inference,
inputs=[gr.Image(type="pil"), gr.Textbox(label="Question")],
outputs=[gr.Text()],
title='Document answer questions'
).launch(debug=True)