jonathanjordan21 commited on
Commit
72a5c8e
1 Parent(s): 95bcc3c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +56 -0
app.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import re
3
+ from PIL import Image
4
+ from io import BytesIO
5
+ import torch
6
+
7
+ from transformers import DonutProcessor, VisionEncoderDecoderModel
8
+
9
+ # Check GPU
10
+ device = "cuda" if torch.cuda.is_available() else "cpu"
11
+
12
+ # Load processor
13
+ processor = DonutProcessor.from_pretrained("jonathanjordan21/donut-finetuned-drugs-composition-indonesian")
14
+
15
+ # Load model
16
+ model = VisionEncoderDecoderModel.from_pretrained("jonathanjordan21/donut-finetuned-drugs-composition-indonesian")
17
+
18
+
19
+ def predict(inp):
20
+ # Define Json Parser
21
+ def get_komposisi(image_path, image=None):
22
+ image = Image.open(image_path).convert('RGB') if image== None else image.convert('RGB')
23
+
24
+ task_prompt = "<s_kmpsi>"
25
+ decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
26
+
27
+ pixel_values = processor(image, return_tensors="pt").pixel_values
28
+
29
+ outputs = model.generate(
30
+ pixel_values.to(device),
31
+ decoder_input_ids=decoder_input_ids.to(device),
32
+ max_length=model.decoder.config.max_position_embeddings,
33
+ early_stopping=True,
34
+ pad_token_id=processor.tokenizer.pad_token_id,
35
+ eos_token_id=processor.tokenizer.eos_token_id,
36
+ use_cache=True,
37
+ bad_words_ids=[[processor.tokenizer.unk_token_id]],
38
+ return_dict_in_generate=True,
39
+ )
40
+
41
+ sequence1 = processor.batch_decode(outputs.sequences)[0]
42
+ sequence2 = sequence1.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
43
+ sequence3 = re.sub(r"<.*?>", "", sequence2, count=1).strip() # remove first task start token
44
+
45
+ return processor.token2json(sequence3)
46
+
47
+ #Generate Output
48
+ out = get_komposisi("", inp)
49
+ return out
50
+
51
+
52
+
53
+ gr.Interface(fn=predict,
54
+ inputs=gr.Image(type="pil"),
55
+ outputs="json").launch()
56
+