RamzesIII commited on
Commit
60b194f
·
1 Parent(s): 8aceff3

create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
2
+ import gradio as gr
3
+
4
+ qa_pipeline = pipeline("document-question-answering", model="RamzesIII/llmv2-docvqa-finetuned")
5
+
6
+ ru_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
7
+ ru_en_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
8
+
9
+ en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
10
+ en_ru_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
11
+
12
+ def translate_ru2en(ru_question, ru_en_model, ru_tokenizer):
13
+ input_ids = ru_tokenizer.encode(ru_question, return_tensors="pt")
14
+ output_ids = ru_en_model.generate(input_ids, max_new_tokens=512)
15
+ en_question = ru_tokenizer.decode(output_ids[0], skip_special_tokens=True)
16
+ return en_question
17
+
18
+
19
+ def translate_en2ru(en_answer, en_ru_model, en_tokenizer):
20
+ input_ids = en_tokenizer.encode(en_answer, return_tensors="pt")
21
+ output_ids = en_ru_model.generate(input_ids, max_new_tokens=512)
22
+ ru_answer = en_tokenizer.decode(output_ids[0], skip_special_tokens=True)
23
+ return ru_answer
24
+
25
+ def ru_inference(image, ru_question):
26
+ en_question = translate_ru2en(ru_question, ru_en_model, ru_tokenizer)
27
+ en_answer = qa_pipeline(image=image, question=en_question)[0]['answer']
28
+ ru_answer = translate_en2ru(en_answer, en_ru_model, en_tokenizer)
29
+ return en_answer
30
+
31
+
32
+ interface = gr.Interface(
33
+ fn=ru_inference,
34
+ inputs=[gr.Image(type="pil"), gr.Textbox(label="Question")],
35
+ outputs=[gr.Text()],
36
+ title='Document answer questions'
37
+ ).launch(debug=True)