S1mohaan commited on
Commit
10bfa1c
1 Parent(s): 84c27be

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
3
+ import requests
4
+ from PIL import Image
5
+ import torch, os, re, json
6
+ import spaces
7
+
8
+ torch.hub.download_url_to_file('https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/test/png/74801584018932.png', 'chart_example_1.png')
9
+ torch.hub.download_url_to_file('https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/val/png/multi_col_1229.png', 'chart_example_2.png')
10
+
11
+
12
+
13
+ model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma")
14
+ processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")
15
+
16
+
17
+ @spaces.GPU
18
+ def predict(image, input_text):
19
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
20
+ model.to(device)
21
+
22
+ image = image.convert("RGB")
23
+
24
+ inputs = processor(text=input_text, images=image, return_tensors="pt")
25
+ inputs = {k: v.to(device) for k, v in inputs.items()}
26
+
27
+ prompt_length = inputs['input_ids'].shape[1]
28
+
29
+ # Generate
30
+ generate_ids = model.generate(**inputs, max_new_tokens=512)
31
+ output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
32
+
33
+ return output_text
34
+
35
+
36
+ image = gr.components.Image(type="pil", label="Chart Image")
37
+ input_prompt = gr.components.Textbox(label="Input Prompt")
38
+ model_output = gr.components.Textbox(label="Model Output")
39
+ examples = [["chart_example_1.png", "Describe the trend of the mortality rates for children before age 5"],
40
+ ["chart_example_2.png", "What is the share of respondants who prefer Facebook Messenger in the 30-59 age group?"]]
41
+
42
+ title = "AI Chart Captioning Bot"
43
+ interface = gr.Interface(fn=predict,
44
+ inputs=[image, input_prompt],
45
+ outputs=model_output,
46
+ examples=examples,
47
+ title=title,
48
+ theme='gradio/soft')
49
+
50
+ interface.launch()