import gradio as gr
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
import requests
from PIL import Image
import torch, os, re, json
import spaces

torch.hub.download_url_to_file('https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/test/png/74801584018932.png', 'chart_example_1.png')
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')



model = PaliGemmaForConditionalGeneration.from_pretrained("ahmed-masry/chartgemma")
processor = AutoProcessor.from_pretrained("ahmed-masry/chartgemma")



prompt_format = """
Given the question, chart analysis output answer and additional context, come up with an appropriate robust answer.

### Question
{}

### Chart Analysis Output
{}

### Context
{}

### Answer
"""


@spaces.GPU
def predict(image, input_text, input_context):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    image = image.convert("RGB")

    inputs = processor(text=input_text, images=image, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    prompt_length = inputs['input_ids'].shape[1]
    
    # Generate chart text
    generate_ids = model.generate(**inputs, max_new_tokens=512)
    output_text = processor.batch_decode(generate_ids[:, prompt_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

    # Generate next layer prompt (WIP)
    chained_prompt = prompt_format.format(input_text, output_text, input_context)

    
    return output_text

   
image = gr.components.Image(type="pil", label="Chart Image")
input_prompt = gr.components.Textbox(label="Input Prompt")
added_context = gr.components.Textbox(label="Input Context")

model_output = gr.components.Textbox(label="Model Output")
examples = [["chart_example_1.png", "Describe the trend of the mortality rates for children before age 5", "The country of interest is Bahrain."],
            ["chart_example_2.png", "What is the share of respondants who prefer Facebook Messenger in the 30-59 age group?", "This data needs to be cleaned"]]

title = "Gradio Demo for ChartGemma + llama3 context"
interface = gr.Interface(fn=predict, 
                         inputs=[image, input_prompt, added_context], 
                         outputs=model_output, 
                         examples=examples, 
                         title=title,
                         theme='gradio/soft')

interface.launch()