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)