Spaces:
Sleeping
Sleeping
Create app.py
Browse files
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()
|