"""
Reference: https://bard.google.com/
"""

import os
import json
import tqdm
import argparse

import sys
sys.path.append('../')
from utilities import *

from models import bard

def verify_response(response):
    if isinstance(response, str):
        response = response.strip() 
    if response == "" or response == None:
        return False
    if "Response Error" in response:
        return False
    return True


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_dir', type=str, default='../data')
    parser.add_argument('--output_dir', type=str, default='../data/texts')
    parser.add_argument('--output_file', type=str, default='captions_bard.json')
    parser.add_argument('--model', type=str, default='bard', help='model name')
    parser.add_argument('--key', type=str, default='', help='key for bard api')
    args = parser.parse_args()

    # data_files = ["testmini.json", "test.json"]
    data_files = ["testmini.json"]

    data = {}
    for data_file in data_files:
        d = json.load(open(os.path.join(args.output_dir, data_file), 'r'))
        data.update(d)

    pids = list(data.keys())
    print("number of images: ", len(pids))

    ### Read skipped image ids
    result_file = os.path.join("../results/bard/output_bard.json")
    results = json.load(open(result_file, 'r'))
    skipped = []
    for pid in results:
        response = results[pid]['response']
        if "I can't help with images of people" in response:
            skipped.append(pid)
    print("number of skipped images: ", len(skipped))

    # output file
    output_file = os.path.join(args.output_dir, args.output_file)
    
    # final test pids
    if os.path.exists(output_file):
        print("Loading existing results...")
        results = json.load(open(output_file, 'r'))
    else:
        results = {
            "model": "bard",
            "url": "https://bard.google.com/",
            "date": "2023-09-13",
            "texts": {}
        }

    test_pids = []
    texts = results['texts']

    for pid in pids:
        if pid in skipped:
            continue
        if pid in texts:
            response = texts[pid]
            if not verify_response(response):
                test_pids.append(pid)
        else:
            test_pids.append(pid)

    print("\nfinal number of images to run: ", len(test_pids))

    # build model
    query = "Describe the fine-grained content of the image or figure, including scenes, objects, relationships, and any text present."
    
    if args.model == 'bard':
        if args.key == '':
            print("Loading key from environment variable")
            key = os.environ['_BARD_API_KEY']
        else:
            key = args.key
        model = bard.Bard_Model(key)

    for pid in tqdm.tqdm(test_pids):
        image_path = os.path.join(args.image_dir, data[pid]["image"])

        print(f"Generating response for {pid}...")
        try:
            response = model.get_response(image_path, query)
            # print(f"Response: {response}")
            results['texts'][pid] = response
            with open(output_file, 'w') as f:
                json.dump(results, f, indent=2)

        except Exception as e:
            print(e)


    
data_prepro/generate_caption_bard.py
lupantech/MathVista
import os
import json
import torch
import warnings

from tqdm import tqdm
from PIL import Image
from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer

warnings.filterwarnings("ignore")


def predict_caption(image_paths):
    images = []
    for image_path in image_paths:
        i_image = Image.open(image_path)
        if i_image.mode != "RGB":
            i_image = i_image.convert(mode="RGB")

        images.append(i_image)

    pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)

    output_ids = model.generate(pixel_values, **gen_kwargs)

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds


if __name__ == "__main__":
    ## Arguments
    input_path = "../data"
    output_path = "../data/texts"
    output_name = "captions_bard.json"

    model_name = "nlpconnect/vit-gpt2-image-captioning"
    feature_extractor_name = "nlpconnect/vit-gpt2-image-captioning"
    tokenizer_name = "nlpconnect/vit-gpt2-image-captioning"

    max_length = 64
    num_beams = 4
    gen_kwargs = {"max_length": max_length, "num_beams": num_beams}

    ## Read data

    data_files = ["testmini.json", "test.json"]
    problems = {}
    for data_file in data_files:
        d = json.load(open(os.path.join(input_path, data_file), 'r'))
        problems.update(d)

    pids = list(problems.keys())

    print("number of images: ", len(pids))
    print(pids[:10])

    ## Prepare the model
    # url = "https://huggingface.co/sachin/vit2distilgpt2"
    model = VisionEncoderDecoderModel.from_pretrained(model_name)
    feature_extractor = ViTFeatureExtractor.from_pretrained(feature_extractor_name)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    ## Generate image captions
    captions = {}

    print(f"Generating captions!")
    for pid in tqdm(pids):
        image_file = os.path.join(input_path, problems[pid]['image'])
        try:
            caption = predict_caption([image_file])[0]
            captions[pid] = caption.capitalize() + '.'
        except Exception as e:
            captions[pid] = ""
            print(image_file)
            print(e)

    ## Save the captions
    output_file = os.path.join(output_path, output_name)
    os.makedirs(output_path, exist_ok=True)
    print(f"Saved to {output_file}")

    results = {
        "model": model_name,
        "feature_extractor": feature_extractor_name,
        "tokenizer": tokenizer_name,
        "max_length": max_length,
        "num_beams": num_beams,
        "texts": captions
    }

    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2, separators=(',', ': '))

data_prepro/generate_caption_vitgpt2.py
lupantech/MathVista
"""
Reference: https://github.com/JaidedAI/EasyOCR 
"""

# ! pip install easyocr

import os
import easyocr
import json
import tqdm

if __name__ == "__main__":

    data_files = ["../data/testmini.json", "../data/test.json"]
    image_dir = "../data"
    output_file = "../data/texts/ocrs_easyocr.json"

    data = {}
    for data_file in data_files:
        d = json.load(open(data_file, 'r'))
        data.update(d)

    pids = list(data.keys())
    print("number of images: ", len(pids))

    reader = easyocr.Reader(['en']) # ocr reader

    results = {
        "model": "easyocr",
        "url": "https://github.com/JaidedAI/EasyOCR",
        "version": "1.1.8",
        "date": "2023-04-06",
        "texts": {}
    }
    for pid in tqdm.tqdm(pids):
        image_file = os.path.join(image_dir, data[pid]["image"])

        try:
            result = reader.readtext(image_file)
        except Exception as e:
            print(image_file, e)
            result = ""
        
        results["texts"][pid] = str(result)
        # print(result)
        # break

    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)
    
data_prepro/generate_text_easyocr.py
lupantech/MathVista


# pids: 799, 681, 615
shot_examples = [
{
"question": "How much money does Ruth need to buy a baking dish, a casserole dish, and an ice cream scoop? (Unit: $)",
"caption": "The image shows a table with a variety of items on it, including a baking dish, ice cream scoop, casserole dish, and rolling pin. The text in the image says:\n\n```\nbaking dish\n$4.00\nice cream scoop\n$6.00\ncasserole dish\n$3.00\nrolling pin\n$4.00\n```",
"ocr": "[([5, 3], 'baking dish'), ([177, 5], '$4.00'), ([7, 41], 'ice cream scoop'), ([177, 37], '$6.00'), ([9, 69], 'casserole dish'), ([177, 69], '$3.00'), ([5, 98], 'rolling pin'), ([177, 101], '$4.00')]",
"solution": """
Find the total cost of a baking dish, a casserole dish, and an ice cream scoop.\n\n$4.00 + $3.00 + $6.00 = $13.00\n\nRuth needs $13.00.
""",
"code": """
baking_dish_price = 4.00
casserole_dish_price = 3.00
ice_cream_scoop_price = 6.00

ans = baking_dish_price + casserole_dish_price + ice_cream_scoop_price
print(ans)
"""
},

{
"question": "What is the largest city in the nation where this plane is headquartered?",
"choices": ['hong kong', 'osaka', 'shanghai', 'tokyo'],
"caption": "The image shows a large passenger jet parked on a tarmac at an airport. The jet is white with red trim and has a red tail. It is sitting on top of a tarmac next to a building. The jet is being loaded with passengers and cargo. The text on the image says \"Japan. Endless Discovery\".",
"solution": """
The caption mentions that the text on the image says "Japan. Endless Discovery". This indicates that the plane is headquartered in Japan. 

Among the Japanese cities, Tokyo is the largest city.

Thus, the answer is D (tokyo).
""",
"code": """
def largest_city(caption, choices):
    countries_largest_cities = {
        'Japan': 'tokyo',
        'China': 'shanghai'
    }

    if "Japan" in caption:
        country = 'Japan'
    elif "China" in caption:
        country = 'China'

    for choice in choices:
        if choice == countries_largest_cities[country]:
            return choice
    return ""

choices = ['hong kong', 'osaka', 'shanghai', 'tokyo']
caption = "The image shows a large passenger jet parked on a tarmac at an airport. The jet is white with red trim and has a red tail. It is sitting on top of a tarmac next to a building. The jet is being loaded with passengers and cargo. The text on the image says 'Japan. Endless Discovery'."

print(largest_city(caption, choices))
"""
},

{
"question": "If two sides of a triangle measure 12 and 7, which of the following cannot be the perimeter of the triangle?",
"choices": ['29', '34', '37', '38'],
"caption": "The image shows a triangle with two sides labeled 7 and 12. The triangle is drawn on a white background. There is no text other than the labels.",
"ocr": "[([70, 74], '7'), ([324, 74], '12')]",
"solution": """
To determine which of the given perimeters cannot be possible for the triangle, we apply the triangle inequality theorem. The sum of any two sides of a triangle must be greater than the third side.

For the maximum possible value of the third side:
12 + 7 = 19

The minimum possible value for the third side:
12 - 7 = 5

The third side for each option:
(A) 29 - 12 - 7 = 10 (valid)
(B) 34 - 12 - 7 = 15 (valid)
(C) 37 - 12 - 7 = 18 (valid)
(D) 38 - 12 - 7 = 19 (invalid because it should be less than 19)

Thus, the answer is D.
""",
"code": """
def is_valid_triangle(a, b, perimeter):
    # Given a and b, find the third side
    third_side = perimeter - a - b
    
    # Check triangle inequality
    if (a + b > third_side) and (a + third_side > b) and (b + third_side > a):
        return True
    return False

# Given sides
a = 12
b = 7

# Given perimeters
perimeters = [29, 34, 37, 38]

# Check which perimeter is not valid
for p in perimeters:
    if not is_valid_triangle(a, b, p):
        print(p)
""",

}
]


def refine_caption(caption):
    if isinstance(caption, str):
        nonsense = ["Sure. ", 
                    "Sure, I can do that.",
                    "Sorry, I can't help with images of people yet.",
                    "I can't process this file.",
                    "I'm unable to help you with that, as I'm only a language model and don't have the necessary information or abilities.",
                    "I'm not programmed to assist with that.",
                    "Please let me know if you have any other questions.",
                    "I hope this is helpful!",
                    "I hope this helps!"]
        for non in nonsense:
            caption = caption.replace(non, "").strip()
        caption = caption.replace("  ", " ").strip()
    else:
        caption = ""
    return caption


def refine_ocr(ocr):
    """
    [   (
        [161, 39], [766, 39], [766, 120], [161, 120]], 
        'The spring force does', 
        0.912845069753024
        ), 
    ]
    -->
    [   (
        [161, 39], 
        'The spring force does', 
        ), 
    ]
    """
    try:
        ocr = eval(ocr)
        if len(ocr) > 0:
            ocr = [([int(e[0][0][0]), int(e[0][0][1])], e[1]) for e in ocr]
            ocr = str(ocr)
        else:
            ocr = ""
    except:
        ocr = ""
    return ocr


def create_one_query(problem, examples, shot_num, shot_type, use_caption, use_ocr):


    ### [1] Demo prompt
    if shot_num == 0:
        demo_prompt = ""
    else:
        demos = []
        shot_num = min(shot_num, len(examples))
        for example in examples[:shot_num]:
            prompt = ""
    
            # question
            prompt += f"Question: {example['question']}"

            # choices
            if "choices" in example:
                texts = ["Choices:"]
                for i, choice in enumerate(example['choices']):
                    texts.append(f"({chr(ord('A')+i)}) {choice}")
                prompt += "\n" + "\n".join(texts)

            # caption
            if use_caption:
                caption = example['caption'] if 'caption' in example else ""
                if caption != "":
                    prompt += "\n" + f"Image description: {caption}"

            # ocr
            if use_ocr:
                ocr = example['ocr'] if 'ocr' in example else ""
                if ocr != "":
                    prompt += "\n" + f"Image detected text: {ocr}"

            # solution
            if shot_type == 'solution':
                solution = example['solution'].strip()
                prompt += "\n" + f"Solution: {solution}"

            # code
            if shot_type == 'code':
                code = example['code'].strip()
                prompt += "\n" + f"Python code: {code}"

            demos.append(prompt)

        demo_prompt = "\n\n".join(demos)

    ### [2] Test query
    # problem info
    question = problem['question']
    unit = problem['unit']
    choices = problem['choices']
    caption = problem['caption']
    ocr = problem['ocr']
    precision = problem['precision']
    question_type = problem['question_type']
    answer_type = problem['answer_type']

    # hint
    if shot_type == 'solution':
        if question_type == "multi_choice":
            assert answer_type == "text"
            hint_text = f"Hint: Please answer the question and provide the correct option letter, e.g., A, B, C, D, at the end."
        else:
            assert answer_type in ["integer", "float", "list"]
            if answer_type == "integer":
                hint_text = f"Hint: Please answer the question requiring an integer answer and provide the final value, e.g., 1, 2, 3, at the end."
            
            elif answer_type == "float" and precision == 1:
                hint_text = f"Hint: Please answer the question requiring a floating-point number with one decimal place and provide the final value, e.g., 1.2, 1.3, 1.4, at the end."
            
            elif answer_type == "float" and precision == 2:
                hint_text = f"Hint: Please answer the question requiring a floating-point number with two decimal places and provide the final value, e.g., 1.23, 1.34, 1.45, at the end."
            
            elif answer_type == "list":
                hint_text = f"Hint: Please answer the question requiring a Python list as an answer and provide the final list, e.g., [1, 2, 3], [1.2, 1.3, 1.4], at the end."
    else:
        assert shot_type == 'code'
        hint_text = "Hint: Please generate a python code to solve the problem"

    # question
    question_text = f"Question: {question}"
    if unit:
        question_text += f" (Unit: {unit})"

    # choices
    if choices:
        # choices: (A) 1.2 (B) 1.3 (C) 1.4 (D) 1.5
        texts = ["Choices:"]
        for i, choice in enumerate(choices):
            texts.append(f"({chr(ord('A')+i)}) {choice}")
        choices_text = "\n".join(texts)
    else:
        choices_text = ""

    # caption
    caption_text = ""
    if use_caption and caption != "":
        caption_text = f"Image description: {caption}"
   
   # ocr
    ocr_text = ""
    if use_ocr and ocr != "":
        ocr_text = f"Image detected text: {ocr}"

    # prompt
    if shot_type == 'solution':
        prompt = "Solution: "
    else:
        assert shot_type == 'code'
        prompt = "Python code: "
    
    elements = [question_text, choices_text, caption_text, ocr_text, hint_text, prompt]
    test_query = "\n".join([e for e in elements if e != ""])

    ### [3] Final query
    query = demo_prompt + "\n\n" + test_query
    query = query.strip()
    return query


def create_query_data(data, caption_data, ocr_data, args):
    query_data = {}

    for pid, problem in data.items():
        if pid in caption_data:
            caption = caption_data[pid]
            caption = refine_caption(caption)
            problem['caption'] = caption
        else:
            problem['caption'] = ""

        if pid in ocr_data:
            ocr = ocr_data[pid]
            ocr = refine_ocr(ocr)
            problem['ocr'] = ocr
        else:
            problem['ocr'] = []

        query = create_one_query(
            problem = problem, 
            examples = shot_examples,
            shot_num = args.shot_num,
            shot_type = args.shot_type,
            use_caption = args.use_caption,
            use_ocr = args.use_ocr
            )
        query_data[pid] = query

    return query_data


if __name__ == "__main__":
    for example in shot_examples:
        print("----------------------------------------")
        print("\nQuestion:", example['question'])
        if "choices" in example:
            print("\nChoices:", example['choices'])
        print("\nCaption:", example['caption'])
        if "ocr" in example:
            print("\nOCR:", example['ocr'])
        print("\nSolution:", example['solution'])
        print("\nCode:", example['code'])
        
        # execute code
        exec(example['code'])
        


evaluation/build_query.py
lupantech/MathVista
import os
import re
import argparse
import pandas as pd

# !pip install python-Levenshtein
from Levenshtein import distance

import sys
sys.path.append('../')
from utilities import *


def get_most_similar(prediction, choices):
    """
    Use the Levenshtein distance (or edit distance) to determine which of the choices is most similar to the given prediction
    """
    distances = [distance(prediction, choice) for choice in choices]
    ind = distances.index(min(distances))
    return choices[ind]
    # return min(choices, key=lambda choice: distance(prediction, choice))


def normalize_extracted_answer(extraction, choices, question_type, answer_type, precision):
    """
    Normalize the extracted answer to match the answer type
    """
    if question_type == 'multi_choice':
        # make sure the extraction is a string
        if isinstance(extraction, str):
            extraction = extraction.strip()
        else:
            try:
                extraction = str(extraction)
            except:
                extraction = ""
    
        # extract "A" from "(A) text"
        letter = re.findall(r'\(([a-zA-Z])\)', extraction)
        if len(letter) > 0:
            extraction = letter[0].upper()
        
        options = [chr(ord('A') + i) for i in range(len(choices))]
            
        if extraction in options:
            # convert option letter to text, e.g. "A" -> "text"
            ind = options.index(extraction)
            extraction = choices[ind]
        else:
            # select the most similar option
            extraction = get_most_similar(extraction, choices)
        assert extraction in choices

    elif answer_type == 'integer':
        try:
            extraction = str(int(float(extraction)))
        except:
            extraction = None

    elif answer_type == 'float':
        try:
            extraction = str(round(float(extraction), precision))
        except:
            extraction = None
        
    elif answer_type == 'list':
        try:
            extraction = str(extraction)
        except:
            extraction = None

    return extraction
    

def safe_equal(prediction, answer):
    """
    Check if the prediction is equal to the answer, even if they are of different types
    """
    try:
        if prediction == answer:
            return True
        return False
    except Exception as e:
        print(e)
        return False


def get_acc_with_contion(res_pd, key, value):
    if key == 'skills':
        # if value in res_pd[key]:
        total_pd = res_pd[res_pd[key].apply(lambda x: value in x)]
    else:
        total_pd = res_pd[res_pd[key] == value]

    correct_pd = total_pd[total_pd['true_false'] == True]
    acc = "{:.2f}".format(len(correct_pd) / len(total_pd) * 100)
    return len(correct_pd), len(total_pd), acc
        
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_dir', type=str, default='../results')
    parser.add_argument('--output_file', type=str, default='output.json')
    parser.add_argument('--score_file', type=str, default='scores.json')
    parser.add_argument('--gt_file', type=str, default='../data/testmini.json', help='ground truth file')
    parser.add_argument('--number', type=int, default=-1, help='number of problems to run')
    parser.add_argument('--rerun', action='store_true', help='rerun the evaluation')
    parser.add_argument('--caculate_gain', action='store_true', help='caculate the socre gains over random guess')
    parser.add_argument('--random_file', type=str, default='score_random_guess.json')  
    args = parser.parse_args()

    # args
    output_file = os.path.join(args.output_dir, args.output_file)

    # # quick test
    # output_file = '../results/llava-llama-2-13b/output_llava_llama_2_13b.json'

    # read json
    print(f"Reading {output_file}...")
    results = read_json(output_file)

    # read ground truth
    print(f"Reading {args.gt_file}...")
    gts = read_json(args.gt_file)

    # full pids
    full_pids = list(results.keys())
    if args.number > 0:
        full_pids = full_pids[:min(args.number, len(full_pids))]
    print("Number of testing problems:", len(full_pids))
    
    ## [1] Evaluate if the prediction is true or false
    print("\nEvaluating the predictions...")
    update_json_flag = False
    for pid in full_pids:
        problem = results[pid]
        # print(problem)

        if args.rerun:
            if 'prediction' in problem:
                del problem['prediction']
            if 'true_false' in problem:
                del problem['true_false']

        choices = problem['choices']
        question_type = problem['question_type']
        answer_type = problem['answer_type']
        precision = problem['precision']
        extraction = problem['extraction']

        if 'answer' in problem:
            answer = problem['answer']
        else:
            answer = gts[pid]['answer']
            problem['answer'] = answer

        # normalize the extracted answer to match the answer type
        prediction = normalize_extracted_answer(extraction, choices, question_type, answer_type, precision)

        # verify the prediction is true or false
        true_false = safe_equal(prediction, answer)
        
        # update the problem
        if "true_false" not in problem:
            update_json_flag = True

        elif true_false != problem['true_false']:
            update_json_flag = True

        if "prediction" not in problem:
            update_json_flag = True

        elif prediction !=  problem['prediction']:
            update_json_flag = True
            
        problem['prediction'] = prediction
        problem['true_false'] = true_false

    # save the updated json
    if update_json_flag:
        print("\n!!!Some problems are updated.!!!")
        print(f"\nSaving {output_file}...")
        save_json(results, output_file)

    ## [2] Calculate the average accuracy
    total = len(full_pids)
    correct = 0
    for pid in full_pids:
        if results[pid]['true_false']:
            correct += 1
    accuracy = str(round(correct / total * 100, 2))
    print(f"\nCorrect: {correct}, Total: {total}, Accuracy: {accuracy}%")

    scores = {"average": {"accuracy": accuracy, "correct": correct, "total": total}}
    
    ## [3] Calculate the fine-grained accuracy scores
    
    # merge the 'metadata' attribute into the data
    for pid in results:
        results[pid].update(results[pid].pop('metadata'))

    # convert the data to a pandas DataFrame
    df = pd.DataFrame(results).T

    print(len(df))
    print("Number of test problems:", len(df))
    # assert len(df) == 1000 # Important!!!

    # asign the target keys for evaluation
    target_keys = ['question_type', 'answer_type', 'language', 'source', 'category', 'task', 'context', 'grade', 'skills']
     
    for key in target_keys:
        print(f"\nType: [{key}]")
        # get the unique values of the key
        if key == 'skills':
            # the value is a list
            values = []
            for i in range(len(df)):
                values += df[key][i]
            values = list(set(values))
        else:
            values = df[key].unique()
        #print(values)

        # calculate the accuracy for each value
        scores[key] = {}
        for value in values:
            correct, total, acc = get_acc_with_contion(df, key, value)
            if total > 0:
                print(f"[{value}]: {acc}% ({correct}/{total})")
                scores[key][value] = {"accuracy": acc, "correct": correct, "total": total}
        
        # sort the scores by accuracy
        scores[key] = dict(sorted(scores[key].items(), key=lambda item: float(item[1]['accuracy']), reverse=True))

    # save the scores
    scores_file = os.path.join(args.output_dir, args.score_file)
    print(f"\nSaving {scores_file}...")
    save_json(scores, scores_file)
    print("\nDone!")

    # [4] Calculate the score gains over random guess
    if args.caculate_gain:
        random_file = os.path.join(args.output_dir, args.random_file)
        random_scores = json.load(open(random_file))

        print("\nCalculating the score gains...")
        for key in scores:
            if key == 'average':
                gain = round(float(scores[key]['accuracy']) - float(random_scores[key]['accuracy']), 2)
                scores[key]['acc_gain'] = gain
            else:
                for sub_key in scores[key]:
                    gain = round(float(scores[key][sub_key]['accuracy']) - float(random_scores[key][sub_key]['accuracy']), 2)
                    scores[key][sub_key]['acc_gain'] = str(gain)

        # save the score gains
        print(f"\nSaving {scores_file}...")    
        save_json(scores, scores_file)
        print("\nDone!")

evaluation/calculate_score.py
lupantech/MathVista
import os
import re
import time
import argparse

from tqdm import tqdm

import sys
sys.path.append('../')
from utilities import *

# OpenAI
import openai
openai.api_key = os.getenv("OPENAI_API_KEY")
# print(openai.api_key)

# load demo prompt
from prompts.ext_ans import demo_prompt


def verify_extraction(extraction):
    extraction = extraction.strip()
    if extraction == "" or extraction == None:
        return False
    return True


def create_test_prompt(demo_prompt, query, response):
    demo_prompt = demo_prompt.strip()
    test_prompt = f"{query}\n\n{response}"
    full_prompt = f"{demo_prompt}\n\n{test_prompt}\n\nExtracted answer: "
    return full_prompt


def extract_answer(response, problem, quick_extract=False):
    question_type = problem['question_type']
    answer_type = problem['answer_type']
    choices = problem['choices']
    query = problem['query']

    if response == "":
        return ""
    
    if question_type == 'multi_choice' and response in choices:
        return response
    
    if answer_type == "integer":
        try:
            extraction = int(response)
            return str(extraction)
        except:
            pass

    if answer_type == "float":
        try:
            extraction = str(float(response))
            return extraction
        except:
            pass

    # quick extraction
    if quick_extract:
        print("Quickly extracting answer...")
        # The answer is "text". -> "text"
        try:
            result = re.search(r'The answer is "(.*)"\.', response)
            if result:
                extraction = result.group(1)
                return extraction
        except:
            pass

    # general extraction
    try:
        full_prompt = create_test_prompt(demo_prompt, query, response)
        extraction = get_chat_response(full_prompt, openai.api_key)
        return extraction
    except Exception as e:
        print(e)
        print(f"Error in extracting answer for {pid}")

    return ""


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # input
    parser.add_argument('--output_dir', type=str, default='../results')
    parser.add_argument('--output_file', type=str, default='answer.json')
    parser.add_argument('--response_label', type=str, default='response', help='response label for the input file')
    # model
    parser.add_argument('--llm_engine', type=str, default='gpt-4-0613', help='llm engine',
                        choices = ['gpt-3.5-turbo', 'gpt-3.5', 'gpt-4', 'gpt-4-0314', 'gpt-4-0613'])
    parser.add_argument('--number', type=int, default=-1, help='number of problems to run')
    parser.add_argument('--quick_extract', action='store_true', help='use rules to extract answer for some problems')
    parser.add_argument('--rerun', action='store_true', help='rerun the answer extraction')
    # output
    parser.add_argument('--save_every', type=int, default=10, help='save every n problems')
    parser.add_argument('--output_label', type=str, default='', help='label for the output file')
    args = parser.parse_args()

    # args
    label = args.response_label
    result_file = os.path.join(args.output_dir, args.output_file)

    if args.output_label != '':
        output_file = result_file.replace('.json', f'_{args.output_label}.json')
    else:
        output_file = result_file

    # read results
    print(f"Reading {result_file}...")
    results = read_json(result_file)

    # full pids
    full_pids = list(results.keys())
    if args.number > 0:
        full_pids = full_pids[:min(args.number, len(full_pids))]
    print("Number of testing problems:", len(full_pids))

    # test pids
    if args.rerun:
        test_pids = full_pids
    else:
        test_pids = []
        for pid in full_pids:
            # print(pid)
            if 'extraction' not in results[pid] or not verify_extraction(results[pid]['extraction']):
                test_pids.append(pid)
    
    test_num = len(test_pids)
    print("Number of problems to run:", test_num)
    # print(test_pids)

    # tqdm, enumerate results
    for i, pid in enumerate(tqdm(test_pids)):
        problem = results[pid]

        assert label in problem
        response = problem[label]       

        
        extraction  = extract_answer(response, problem, args.quick_extract)
        results[pid]['extraction'] = extraction

        if i % args.save_every == 0 or i == test_num - 1:
            print(f"Saving results to {output_file}...")
            save_json(results, output_file)
            print(f"Results saved.")

evaluation/extract_answer.py
lupantech/MathVista
import os
import io
import time
import argparse

from tqdm import tqdm

import sys
sys.path.append('../')
from utilities import *

from models import claude, gpt, bard

from build_query import create_query_data


def verify_response(response):
    if isinstance(response, str):
        response = response.strip() 
    if response == "" or response == None:
        return False
    if "Response Error" in response:
        return False
    return True


def evaluate_code(code_string):
    # execute_code_and_capture_output
    # Backup the original stdout
    old_stdout = sys.stdout
    
    # Redirect stdout to capture the output
    new_stdout = io.StringIO()
    sys.stdout = new_stdout
    
    # Try executing the code and capture any exception
    error = None
    try:
        exec(code_string)
    except Exception as e:
        error = e
    
    # Restore the original stdout
    sys.stdout = old_stdout
    
    # Get the captured output
    captured_output = new_stdout.getvalue()
    if isinstance(captured_output, str):
        captured_output = captured_output.strip()
    
    # Return the captured output or error
    return captured_output, error
    


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # input
    parser.add_argument('--data_dir', type=str, default='../data')
    parser.add_argument('--input_file', type=str, default='testmini.json')
    # output
    parser.add_argument('--output_dir', type=str, default='../results/bard')
    parser.add_argument('--output_file', type=str, default='output_bard.json')
    # model
    parser.add_argument('--model', type=str, default='gpt-3.5-turbo', help='llm engine',
                        choices = ['gpt-3.5-turbo', 'claude-2', 'gpt4', 'gpt-4-0613', 'bard'])
    parser.add_argument('--key', type=str, default='', help='key for llm api')
    # query
    parser.add_argument('--query_file', type=str, default=None)  
    parser.add_argument('--caption_file', type=str, default='../data/texts/captions_bard.json') 
    parser.add_argument('--ocr_file', type=str, default='../data/texts/ocrs_easyocr.json')   
    parser.add_argument('--shot_type', type=str, default='solution', help='shot type', 
                        choices = ['solution', 'code'])
    parser.add_argument('--shot_num', type=int, default=0, help='number of shot examples')
    parser.add_argument('--use_caption', action='store_true', help='use caption data')
    parser.add_argument('--use_ocr', action='store_true', help='use ocr data')
    # other settings
    parser.add_argument('--rerun', action='store_true', help='rerun answer extraction for all problems')
    parser.add_argument('--debug', action='store_true', help='debug mode')
    args = parser.parse_args()

    # load data
    input_file = os.path.join(args.data_dir, args.input_file)
    print(f"Reading {input_file}...")
    data = read_json(input_file)
    
    # load or create query data
    if args.query_file:
        query_file = os.path.join(args.data_dir, args.query_file)
        if os.path.exists(query_file):
            print(f"Loading existing {query_file}...")
            query_data = read_json(query_file)
    else:
        print("\nCreating new query...")
        # load caption
        caption_data = {}
        if args.use_caption:
            caption_file = args.caption_file
            if os.path.exists(caption_file):
                print(f"Reading {caption_file}...")
                try:
                    caption_data = read_json(caption_file)["texts"]
                    print("Caption data loaded.")
                except:
                    print("Caption data not found!! Please Check.")                    
        # load ocr
        ocr_data = {}
        if args.use_ocr:
            ocr_file = args.ocr_file
            if os.path.exists(ocr_file):
                print(f"Reading {ocr_file}...")
                try:
                    ocr_data = read_json(ocr_file)["texts"]
                    print("OCR data loaded.")
                except:
                    print("OCR data not found!! Please Check.")
        # create query
        query_data = create_query_data(data, caption_data, ocr_data, args)

    # output file
    os.makedirs(args.output_dir, exist_ok=True)
    output_file = os.path.join(args.output_dir, args.output_file)
    
    # load results
    if os.path.exists(output_file):
        print("\nResults already exist.")
        print(f"Reading {output_file}...")
        results = read_json(output_file)
    else:
        results = {}

    # load model
    print(f"\nLoading {args.model}...")
    if args.model == 'bard':
        if args.key == '':
            print("Loading key from environment variable")
            key = os.environ['_BARD_API_KEY']
        else:
            key = args.key
        model = bard.Bard_Model(key)
    
    elif "gpt" in args.model:
        if args.key == '':
            print("Loading token from environment variable")
            key = os.getenv("OPENAI_API_KEY")
        else:
            key = args.key
        model = gpt.GPT_Model(args.model, key)
    
    elif "claude" in args.model:
        if args.key == '':
            print("Loading token from environment variable")
            key = os.environ.get("ANTHROPIC_API_KEY")
        else:
            key = args.key
        model = claude.Claude_Model(args.model, key)
    
    print(f"Model loaded.")
    
    # build final test pid list
    test_pids = list(data.keys())
    print("\nNumber of test problems in total:", len(test_pids))

    skip_pids = []
    if not args.rerun:
        print("\nRemoving problems with existing valid response...")
        for pid in test_pids:
            # print(f"Checking {pid}...")
            if pid in results and 'response' in results[pid]:
                response = results[pid]['response']
                if verify_response(response):
                    # print(f"Valid response found for {pid}.")
                    skip_pids.append(pid)
    else:
        print("\nRerun answer extraction for all problems...")

    test_pids = [pid for pid in test_pids if pid not in skip_pids]
    print("Number of test problems to run:", len(test_pids))
    # print(test_pids)

    # tqdm, enumerate results
    for _, pid in enumerate(tqdm(test_pids)):
        problem = data[pid]
        query = query_data[pid]
        image = problem['image']
        image_path = os.path.join(args.data_dir, image)

        if args.debug:
            print("--------------------------------------------------------------")
        print(f"\nGenerating response for {pid}...")
        try:
            response = model.get_response(image_path, query)
            # print(f"Response: {response}")
            results[pid] = problem
            results[pid]['query'] = query
            if args.shot_type == 'solution':
                results[pid]['response'] = response
            else:
                output, error = evaluate_code(response)
                results[pid]['response'] = response
                results[pid]['execution'] = output
                results[pid]['error'] = str(error)
            if args.debug:
                print(f"\n#Query: \n{query}")
                print(f"\n#Response: \n{response}")
        except Exception as e:
            print(e)
            print(f"Error in extracting answer for {pid}")
            results[pid]['error'] = e
    
        try:
            print(f"Saving results to {output_file}...")
            save_json(results, output_file)
            print(f"Results saved.")
        except Exception as e:
            print(e)
            print(f"Error in saving {output_file}")

evaluation/generate_response.py
lupantech/MathVista


# pids = 852,  104,  824,  506,  540

demo_prompt = """
Please read the following example. Then extract the answer from the model response and type it at the end of the prompt.

Hint: Please answer the question requiring an integer answer and provide the final value, e.g., 1, 2, 3, at the end.
Question: Which number is missing?

Model response: The number missing in the sequence is 14.

Extracted answer: 14

Hint: Please answer the question requiring a floating-point number with one decimal place and provide the final value, e.g., 1.2, 1.3, 1.4, at the end.
Question: What is the fraction of females facing the camera?

Model response: The fraction of females facing the camera is 0.6, which means that six out of ten females in the group are facing the camera.

Extracted answer: 0.6

Hint: Please answer the question requiring a floating-point number with two decimal places and provide the final value, e.g., 1.23, 1.34, 1.45, at the end.
Question: How much money does Luca need to buy a sour apple candy and a butterscotch candy? (Unit: $)

Model response: Luca needs $1.45 to buy a sour apple candy and a butterscotch candy.

Extracted answer: 1.45

Hint: Please answer the question requiring a Python list as an answer and provide the final list, e.g., [1, 2, 3], [1.2, 1.3, 1.4], at the end.
Question: Between which two years does the line  graph saw its maximum peak?

Model response: The line graph saw its maximum peak between 2007 and 2008.

Extracted answer: [2007, 2008]

Hint: Please answer the question and provide the correct option letter, e.g., A, B, C, D, at the end.
Question: What fraction of the shape is blue?\nChoices:\n(A) 3/11\n(B) 8/11\n(C) 6/11\n(D) 3/5

Model response: The correct answer is (B) 8/11.

Extracted answer: B
"""
evaluation/prompts/ext_ans.py
lupantech/MathVista
import os
import time

# !pip install bardapi
from bardapi import Bard


# verify response
def verify_response(response):
    if isinstance(response, str):
        response = response.strip() 
    if response == "" or response == None:
        return False
    if "Response Error" in response:
        # print("Response Error")
        return False
    return True


# build bard class
class Bard_Model():
    def __init__(self, token, patience=1, sleep_time=1):
        print("!!Bard Token:", token)
        self.patience = patience
        self.sleep_time = sleep_time
        self.model = Bard(token=token)

    def get_response(self, image_path, input_text):
        patience = self.patience
        while patience > 0:
            patience -= 1
            # print(f"Patience: {patience}")
            try:
                assert os.path.exists(image_path)
                image = open(image_path, 'rb').read() # (jpeg, png, webp) are supported.
                response = self.model.ask_about_image(input_text, image)
                response = response['content'].strip()
                if verify_response(response):
                    # print("# Verified Response")
                    # print(response)
                    return response
                else:
                    print(response)
            except Exception as e:
                print(e)
                if self.sleep_time > 0:
                    time.sleep(self.sleep_time)
        return ""

models/bard.py
lupantech/MathVista

import time
from anthropic import Anthropic, HUMAN_PROMPT, AI_PROMPT


# build claude class
class Claude_Model():
    def __init__(self, model="claude-2", api_key="", temperature=0, max_tokens=1024, n=1, patience=1000, sleep_time=0):
        self.model = model
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.n = n
        self.patience = patience
        self.sleep_time = sleep_time

    def get_response(self, image_path, user_prompt):
        patience = self.patience
        while patience > 0:
            patience -= 1
            try:
                # configure the default for all requests:
                anthropic = Anthropic(
                    max_retries=0,
                    api_key=self.api_key,
                )

                # update prompt
                if "python" in user_prompt:
                    _HUMAN_PROMPT = HUMAN_PROMPT + "Generate the runnable python code only."
                else:
                    _HUMAN_PROMPT = HUMAN_PROMPT

                # configure per-request options
                completion = anthropic.with_options(max_retries=5).completions.create(
                    prompt=f"{_HUMAN_PROMPT} {user_prompt}{AI_PROMPT}",
                    max_tokens_to_sample=self.max_tokens,
                    model=self.model,
                )
                
                # inference
                prediction = completion.completion.strip()
                if "python" in user_prompt:
                    prediction = prediction.replace("```python", "").replace("```", "").strip()
                if prediction != "" and prediction != None:
                    return prediction
                        
            except Exception as e:
                if "limit" not in str(e):
                    print(e)
                if self.sleep_time > 0:
                    time.sleep(self.sleep_time)
        return ""

models/claude.py
lupantech/MathVista
import time
import openai


# build gpt class
class GPT_Model():
    def __init__(self, model="gpt-3.5-turbo", api_key="", temperature=0, max_tokens=1024, n=1, patience=1000000, sleep_time=0):
        self.model = model
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.n = n
        self.patience = patience
        self.sleep_time = sleep_time

    def get_response(self, image_path, user_prompt):
        patience = self.patience
        max_tokens = self.max_tokens
        messages = [
            {"role": "user", "content": user_prompt},
        ]
        while patience > 0:
            patience -= 1
            try:
                # print("self.model", self.model)
                response = openai.ChatCompletion.create(model=self.model,
                                                        messages=messages,
                                                        api_key=self.api_key,
                                                        temperature=self.temperature,
                                                        max_tokens=max_tokens,
                                                        n=self.n
                                                        )
                if self.n == 1:
                    prediction = response['choices'][0]['message']['content'].strip()
                    if prediction != "" and prediction != None:
                        return prediction
                else:
                    prediction = [choice['message']['content'].strip() for choice in response['choices']]
                    if prediction[0] != "" and prediction[0] != None:
                        return prediction
                        
            except Exception as e:
                if "limit" not in str(e):
                    print(e)
                if "Please reduce the length of the messages or completion" in str(e):
                    max_tokens = int(max_tokens * 0.9)
                    print("!!Reduce max_tokens to", max_tokens)
                if max_tokens < 8:
                    return ""
                if "Please reduce the length of the messages." in str(e):
                    print("!!Reduce user_prompt to", user_prompt[:-1])
                    return ""
                if self.sleep_time > 0:
                    time.sleep(self.sleep_time)
        return ""

models/gpt.py
lupantech/MathVista
import os
import cv2
import json
import time
import pickle
import openai
import re
from word2number import w2n


def create_dir(output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    

def read_csv(file):
    data = []
    with open(file, 'r') as f:
        for line in f:
            data.append(line.strip())
    return data


def read_pandas_csv(csv_path):
    # read a pandas csv sheet 
    import pandas as pd
    df = pd.read_csv(csv_path)
    return df


def read_json(path):
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


def read_jsonl(file):
    with open(file, 'r') as f:
        data = [json.loads(line) for line in f]
    return data


def read_pickle(path):
    with open(path, 'rb') as f:
        return pickle.load(f)


def save_json(data, path):
    with open(path, 'w') as f:
        json.dump(data, f, indent=4)


def save_array_img(path, image):
    cv2.imwrite(path, image)


def contains_digit(text):
    # check if text contains a digit
    if any(char.isdigit() for char in text):
        return True
    return False  
    
def contains_number_word(text):
    # check if text contains a number word
    ignore_words = ["a", "an", "point"]
    words = re.findall(r'\b\w+\b', text)  # This regex pattern matches any word in the text
    for word in words:
        if word in ignore_words:
            continue
        try:
            w2n.word_to_num(word)
            return True  # If the word can be converted to a number, return True
        except ValueError:
            continue  # If the word can't be converted to a number, continue with the next word
    
    # check if text contains a digit
    if any(char.isdigit() for char in text):
        return True

    return False  # If none of the words could be converted to a number, return False


def contains_quantity_word(text, special_keep_words=[]):
    # check if text contains a quantity word
    quantity_words = ["most", "least", "fewest"
                      "more", "less", "fewer", 
                      "largest", "smallest", "greatest", 
                      "larger", "smaller", "greater", 
                      "highest", "lowest", "higher", "lower",
                      "increase", "decrease",
                      "minimum", "maximum", "max", "min",
                      "mean", "average", "median",
                      "total", "sum", "add", "subtract",
                      "difference", "quotient", "gap",
                      "half", "double", "twice", "triple",
                      "square", "cube", "root",
                      "approximate", "approximation",
                      "triangle", "rectangle", "circle", "square", "cube", "sphere", "cylinder", "cone", "pyramid",
                      "multiply", "divide",
                      "percentage", "percent", "ratio", "proportion", "fraction", "rate", 
                    ]
    
    quantity_words += special_keep_words # dataset specific words
    
    words = re.findall(r'\b\w+\b', text)  # This regex pattern matches any word in the text
    if any(word in quantity_words for word in words):
        return True

    return False  # If none of the words could be converted to a number, return False


def is_bool_word(text):
    if text in ["Yes", "No", "True", "False", 
                "yes", "no", "true", "false", 
                "YES", "NO", "TRUE", "FALSE"]:
        return True
    return False


def is_digit_string(text):
    # remove ".0000"
    text = text.strip()
    text = re.sub(r'\.0+$', '', text)
    try:
        int(text)
        return True
    except ValueError:
        return False
   
    
def is_float_string(text):
    # text is a float string if it contains a "." and can be converted to a float
    if "." in text:
        try:
            float(text)
            return True
        except ValueError:
            return False
    return False


def copy_image(image_path, output_image_path):
    from shutil import copyfile
    copyfile(image_path, output_image_path)


def copy_dir(src_dir, dst_dir):
    from shutil import copytree
    # copy the source directory to the target directory
    copytree(src_dir, dst_dir)


import PIL.Image as Image
def get_image_size(img_path):
    img = Image.open(img_path)
    width, height = img.size
    return width, height


def get_chat_response(promot, api_key, model="gpt-3.5-turbo", temperature=0, max_tokens=256, n=1, patience=10000000,
 sleep_time=0):
    messages = [
        {"role": "user", "content": promot},
    ]
    # print("I am here")
    while patience > 0:
        patience -= 1
        try:
            response = openai.ChatCompletion.create(model=model,
                                                messages=messages,
                                                api_key=api_key,
                                                temperature=temperature,
                                                max_tokens=max_tokens,
                                                n=n)
            if n == 1:
                prediction = response['choices'][0]['message']['content'].strip()
                if prediction != "" and prediction != None:
                    return prediction
            else:
                prediction = [choice['message']['content'].strip() for choice in response['choices']]
                if prediction[0] != "" and prediction[0] != None:
                    return prediction

        except Exception as e:
            if "Rate limit" not in str(e):
                print(e)

            if "Please reduce the length of the messages" in str(e):
                print("!!Reduce promot size")
                # reduce input prompt and keep the tail
                new_size = int(len(promot) * 0.9)
                new_start = len(promot) - new_size
                promot = promot[new_start:]
                messages = [
                    {"role": "user", "content": promot},
                ]
                
            if sleep_time > 0:
                time.sleep(sleep_time)
    return ""

utilities.py
lupantech/MathVista