Spaces:
Sleeping
Sleeping
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) |