preprocesssors.py

import hashlib
import json
import math
from functools import reduce
from typing import Mapping, Optional, Sequence

import numpy as np
import tensorflow as tf
import seqio
import gin

from .data_utils import flatten_parts, stateless_permutation, stateless_shuffle
from .. import config


def get_from_dict(data, keys):
    """Iterate nested dictionary"""
    return reduce(dict.get, keys, data)

def get_blank_image():
    image = tf.zeros([224, 224, 3], dtype=tf.uint8)
    image = tf.expand_dims(image, 0)[:1]
    return image


@seqio.utils.map_over_dataset
def rekey(x, key_map=None):
    """Replace the feature keys according to the mapping in `key_map`.
    For example, if the dataset returns examples of the format:
    {'foo': 'something', 'bar': 'something else'}
    and key_map = {'boo': 'foo', 'spar': 'bar'} then this function will return
    examples with the format
    {'boo': 'something', 'spar': 'something else'}
    If a mapping is to an empty key or None, set the new key to an empty string.
    Args:
        x: an example to process.
        key_map: dictionary mapping new keys to original keys
    Returns:
        A preprocessed example with the format listed above.
    """
    if key_map:
        out = {}
        for new_key, old_key in key_map.items():
            if isinstance(old_key, list):
                out[new_key] = get_from_dict(x, old_key)
            else:
                out[new_key] = x[old_key]
        return out
    return x


def rename(**kwargs):
    @seqio.map_over_dataset
    def _fn(x):
        updates = {}
        for new_key, old_key in kwargs.items():
            if isinstance(old_key, list):
                val = x[old_key[0]]
                for k in old_key[1:-1]:
                    val = val[k]
                updates[new_key] = val.pop(old_key[-1])
            else:
                updates[new_key] = x.pop(old_key)
        x.update(updates)
        return x
    return _fn


def extract_transcripts(ds):
    ds = flatten_parts(ds, ["transcripts"])
    def _map(ex):
        return dict(
            image=ex["image"],
            text=ex["transcripts"],
            url=ex["url"]
        )
    return ds.map(_map)


@seqio.map_over_dataset
def extract_caption_and_all_transcripts(ex):
    transcripts = tf.random.shuffle(ex["transcripts"])[:3]
    weight = 1.0 / tf.cast(tf.shape(transcripts)[0], tf.float32)
    return dict(
        image=ex["image"],
        text=tf.concat([tf.expand_dims(ex["caption"], 0), transcripts], 0),
        url=ex["url"],
        text_weights=tf.pad(
            tf.ones((1,), dtype=tf.float32), [[0, tf.shape(transcripts)[0]]],
            constant_values=weight),
    )


@seqio.map_over_dataset
def extract_all_transcripts(ex):
    transcripts = tf.random.shuffle(ex["transcripts"])[:3]
    weight = 3.0 / tf.cast(tf.shape(transcripts)[0], tf.float32)
    return dict(
        image=ex["image"],
        text=transcripts,
        url=ex["url"],
        text_weights=tf.fill((tf.shape(transcripts)[0],), weight),
    )


@seqio.map_over_dataset
def extract_transcript(ex):
    transcripts = tf.random.shuffle(ex["transcripts"])
    return dict(
        image=ex["image"],
        text=transcripts[0],
        url=ex["url"],
    )


@seqio.map_over_dataset
def extract_caption(ex):
    caption = ex["caption"]
    if len(caption.shape) > 0:
        ex["text"] = caption[0]
    else:
        ex["text"] = caption
    return ex


@seqio.map_over_dataset
def extract_joint_captions(ex):
    caption = ex["caption"]
    if len(caption.shape) > 0:
        caption = caption[0]
    _ix = tf.random.uniform((), 0, tf.shape(ex["transcripts"])[0], dtype=tf.int32)
    _ix = _ix % tf.shape(ex["transcripts"])[0]
    return dict(
        image=ex["image"],
        text=tf.stack([caption, ex["mistral_caption"], ex["transcripts"][_ix]], 0),
        url=ex["url"]
    )


@seqio.map_over_dataset(num_seeds=1)
def extract_caption_and_transcript(ex, seed):
    caption = ex["caption"]
    if len(caption.shape) > 0:
        caption = caption[0]
    _ix = tf.random.stateless_uniform((), seed, 0, tf.shape(ex["transcripts"])[0], dtype=tf.int32)
    return dict(
        image=ex["image"],
        text=tf.stack([caption, ex["transcripts"][_ix]], 0),
        url=ex["url"]
    )


@seqio.map_over_dataset
def caption_transcript_augmented(ex, sequence_length):
    caption = ex["caption"]
    if len(caption.shape) > 0:
        caption = caption[0]
    image = ex["image"]
    properties = []

    do_augmentation = sequence_length["is_training"]
    # do_augmentation = False

    # Keep this off, it screws up OCR
    # do_hflip = (tf.random.uniform(()) > 0.2 and do_augmentation)
    do_hflip = False
    if do_hflip:
        image = image[:, ::-1]

    # Mild color jitter
    do_color = (tf.random.uniform(()) > 0.5 and do_augmentation)
    if do_color:
        image = tf.image.random_hue(image, max_delta=0.05)
        image = tf.image.random_brightness(image, max_delta=0.2)
        image = tf.image.random_saturation(image, 0.7, 1.3)
        image = tf.image.random_contrast(image, 0.7, 1.3)

    # Mild affine transformation
    do_affine = (tf.random.uniform(()) > 0.5 and do_augmentation)
    if do_affine and do_augmentation:
        shift_x = tf.random.uniform((), -10, 10) * 0
        shift_y = tf.random.uniform((), -10, 10) * 0
        shear_x = tf.random.uniform((), -2, 2)
        shear_y = tf.random.uniform((), -2, 2)
        rotation = tf.random.uniform((), -6, 6)
        max_scale = 1.1
        scale = tf.random.uniform((), 0.8, max_scale)
        center = tf.cast(tf.shape(image), tf.float32)/2

        image = tf.keras.ops.image.affine_transform(
            image,
            tf.stack(get_affine_matrix(
                [center[0], center[1]],
                rotation,
                [shift_x, shift_y],
                1/scale,
                [shear_x, shear_y]
            ) + [0., 0.]),
            interpolation='bilinear',
            fill_mode='constant',
            fill_value=1.,
            data_format='channels_last'
        )

    properties = tf.stack([
        ("[hflip]" if do_hflip else ""),
        ("[color]" if do_color else ""),
        ("[affine]" if do_affine else "")
    ])
    properties = tf.boolean_mask(properties, tf.strings.length(properties) > 0)
    prompt = tf.strings.reduce_join(properties, separator=" ")
    ix = tf.random.uniform((), 0, tf.shape(ex["transcripts"])[0], dtype=tf.int32)
    out = dict(
        image=image,
        text=tf.stack([caption, ex["transcripts"][ix]], 0),
        url=ex["url"],
        prompt=prompt,
    )
    # out["metadata/unaugmented_image"] = image
    return out


def extract_caption_and_transcript_hflip(ds):

    # Just in case they are ordered somehow in Matt's data
    @seqio.map_over_dataset
    def _shuffle_transcripts(_ex):
        _ex["transcripts"] = tf.random.shuffle(_ex["transcripts"])
        _ex["hflip"] = tf.random.uniform((), 0, 3, dtype=tf.int32)
        return _ex

    ds = _shuffle_transcripts(ds)

    # Build a 3x long dataset with each individual transcript so we iterate through
    # each transcript
    @seqio.map_over_dataset
    def _with_transcript(ex, _ix):
        caption = ex["caption"]
        if len(caption.shape) > 0:
            caption = caption[0]
        hflip = ex["hflip"] == _ix
        if hflip:
            ex["image"] = ex["image"][:, ::-1]
            style = ["long_caption_flipped", "transcript_flipped"]
        else:
            style = ["long_caption", "transcript"]
        return dict(
            image=ex["image"],
            text=tf.stack([caption, ex["transcripts"][_ix]], 0),
            url=ex["url"],
            style=style
        )

    joint_ds = _with_transcript(ds, 0)
    for i in range(1, 3):
        joint_ds = joint_ds.concatenate(_with_transcript(ds, i))

    return joint_ds


@seqio.map_over_dataset
def extract_llava(ex, sequence_length, output_features):
    tf.assert_equal(tf.shape(ex['conversations']['value'])[0], 2)
    prompt = ex['conversations']['value'][0]
    text = ex['conversations']['value'][1]
    ex.pop('conversations')
    ex["text"] = text
    ex["prompt"] = prompt
    return ex


def extract_localized_narrative(ds):
    ds = ds.filter(lambda ex: tf.shape(ex["cap/cap_caption"])[0] > 0)
    def _map(ex):
        return dict(
            image=ex["image"],
            text=tf.strings.reduce_join(ex["cap/cap_caption"], separator="\n")
        )
    return ds.map(_map)


def float_to_text(val):
    return tf.strings.as_string(tf.cast(val * 100, tf.int32))


@seqio.map_over_dataset
def extract_vqa(ex):
    questions = ex["vqa"]["questions"]
    answers = ex["vqa"]["answers"]
    answers = tf.strings.reduce_join(answers, 1, separator="; ")
    qas = tf.strings.reduce_join(tf.stack([questions, answers], 1), separator=" ")
    return dict(
        image=ex["image"],
        text=tf.strings.reduce_join(qas, separator="\n")
    )


@seqio.map_over_dataset
def coco_image_id_from_path(ex):
    image_id = tf.strings.substr(ex["image/filename"], 0, tf.strings.length(ex["image/filename"])-4)
    ex["image_id"] = tf.strings.to_number(image_id)
    return ex


@seqio.map_over_dataset
def add_coco_url(ex):
    """Turns a COCO path into a URL, which can then be used in visualizations"""
    path = ex["image/filename"]
    if not tf.strings.regex_full_match(path, ".*/.*"):
        prefix = tf.strings.regex_replace(path, "COCO_", "")
        prefix = tf.strings.regex_replace(prefix, "_[0-9]+.jpg", "")
        path = tf.strings.join([prefix, path], separator="/")

    # images are hosted by the COCO website here
    url = tf.strings.join(["https://s3.us-east-1.amazonaws.com/images.cocodataset.org/", path])
    ex["metadata/image_url"] = url
    return ex


def flatten_vqa(ds):
    parts = ["questions", "answers"]
    for k in ["id", "question_id"]:
        if k in ds.element_spec:
            parts.append(k)
    return flatten_parts(ds, parts)


def format_gqa(ds, is_balanced=True, flatten=True):
    if is_balanced:
        ds = ds.filter(lambda x: tf.reduce_any(x["questions"]["is_balanced"]))
        def _filter_qs(ex):
            qs = ex["questions"]
            mask = qs["is_balanced"]
            qs = {k: tf.boolean_mask(v, mask) for k, v in qs.items()}
            ex["questions"] = qs
            return ex
        ds = ds.map(_filter_qs)

    if flatten:
        ds = flatten_parts(ds, ["questions"])

    def _rename(ex):
        out = ex["questions"]
        out["image"] = ex["image"]
        out["image_id"] = ex["image_id"]
        return out
    return ds.map(_rename)


@seqio.map_over_dataset
def fix_doqa_url(x):
    x["image_url"] = tf.strings.regex_replace(x["image_url"], "gs://", "")
    return x


def _add_metadata(ex):
    out = {}
    if "id" in ex:
        out["metadata/example_id"] = ex["id"]
    elif "example_id" in ex:
        out["metadata/example_id"] = ex["example_id"]
    elif "question_id" in ex:
        out["metadata/example_id"] = ex["question_id"]
    if "image_url" in ex:
        out["metadata/image_url"] = ex["image_url"]
    for k, v in ex.items():
        if k.startswith("metadata/"):
            out[k] = v
    return out


def image_only(ds):
    return ds.filter(lambda x: x["has_image"])


def filter_difficult_direct_answer(ds):
    return ds.filter(lambda x: not x["difficult_direct_answer"])


@seqio.map_over_dataset()
def format_ai2d(ex, variable_style=True):
    abc = tf.constant(list("abcdefg".upper()))
    out = dict(image=ex["image"])
    out.update(_add_metadata(ex))

    options = ex["choices"]
    # >= 3 in case of none of the above like answers
    n_options = tf.shape(ex["option_is_abc"])[0]
    if ex["abc_label"] and tf.reduce_sum(tf.cast(ex["option_is_abc"], tf.int32)) >= (n_options - 1):
        # The image labels are always upper, so use upper in the answer ptions
        options = tf.where(
            ex["option_is_abc"],
            tf.strings.upper(options),
            options
        )
        short_options = options
        style = "ai2_diagram_no_letter"
    else:
        short_options = abc[:tf.shape(options)[0]]
        options = tf.stack([short_options, options,], 1)
        options = tf.strings.reduce_join(options, axis=-1, separator=": ")
        style = "ai2_diagram"

    options = tf.strings.reduce_join(options, separator="\n")
    out["question"] = ex["question"]
    out["options"] = options
    if variable_style:
        out["style"] = style
    if ex["answer_idx"] < 0:
        out["text"] = "?"
    else:
        out["text"] = short_options[ex["answer_idx"]]
    out["metadata/answer_idx"] = ex["answer_idx"]
    tf.debugging.assert_equal(tf.reduce_any(tf.strings.regex_full_match(options, ".*\|\|\|.*")), False)
    out["metadata/option_names"] = tf.strings.reduce_join(short_options, separator="|||")
    out["metadata/has_transparent_box"] = ex.get("has_transparent_box", tf.constant(False))
    out["metadata/abc_label"] = ex["abc_label"]
    return out


@gin.configurable()
@seqio.map_over_dataset()
def format_multiple_choice_qa(ex, option_format="abc"):
    assert option_format == "abc"
    abc = tf.constant(list("abcdefg".upper()))
    out = dict(image=ex["image"])
    out.update(_add_metadata(ex))
    options = ex["choices"]
    short_options = abc[:tf.shape(options)[0]]
    options = tf.stack([short_options, options,], 1)
    options = tf.strings.reduce_join(options, axis=-1, separator=": ")
    options = tf.strings.reduce_join(options, separator="\n")
    out["question"] = ex["question"]
    out["options"] = options
    if ex["answer_idx"] < 0:
        out["text"] = "?"
    else:
        out["text"] = short_options[ex["answer_idx"]]
    out["metadata/answer_idx"] = ex["answer_idx"]
    tf.debugging.assert_equal(tf.reduce_any(tf.strings.regex_full_match(options, ".*\|\|\|.*")), False)
    out["metadata/option_names"] = tf.strings.reduce_join(short_options, separator="|||")
    # out["metadata/option_names"] = tf.RaggedTensor.from_row_lengths(short_options, tf.shape(short_options))
    # out["metadata/option_names"] = short_options
    return out


@seqio.map_over_dataset()
def output_options(ex):
    ex["metadata/options"] = ex["options"]
    return ex


@seqio.map_over_dataset()
def extract_tally_qa(ex):
    questions = ex.pop("questions")
    ex["questions"] = questions["question"]
    ex["answers"] = tf.strings.as_string(questions["answer"])
    ex["question_id"] = questions["question_id"]
    return ex


@seqio.map_over_dataset()
def count_bench_preprocessor(ex):
    return {
        "image": ex["image"],
        "text": tf.strings.as_string(ex["number"]),
        "object": ex["noun"],
        "question": tf.strings.join([
            "How many ", ex["noun"], " are there?"
        ]),
        "metadata/count": ex["number"],
    }


def filter_human(ds):
    return ds.filter(lambda x: x["is_human"])


def filter_aug(ds):
    return ds.filter(lambda x: not x["is_human"])


@seqio.map_over_dataset()
def reweight_chartqa(ex, human, aug):
    is_human = ex["metadata/is_human"]
    ex["text_weights"] = human if is_human else aug
    return ex


@seqio.map_over_dataset()
def chartqa_prompting(ex):
    question = tf.strings.join([ex["question"], " Answer:"])
    return dict(
        image=ex["image"],
        question=question,
        answer=ex["answer"]
    )


@seqio.map_over_dataset()
def chartqa_explanation(ex):
    question = tf.strings.join([ex["question"], " Explanation:"])
    out = {
        "image": ex["image"],
        "question": question,
        "answer": ex["answer"],
    }
    out.update({k: v for k, v in ex.items() if k.startswith("metadata/")})
    return out


@seqio.map_over_dataset(num_seeds=1)
def _preprocess_scifi(ex, seed):
    if "qa_pairs" in ex:
        q = ex["qa_pairs"]
    else:
        q = ex["qa"]
    ix = stateless_permutation(tf.shape(q["question"])[0], seed)
    return dict(
        image=ex["image"],
        question=tf.gather(q["question"], ix),
        explanation=tf.gather(q["explanation"], ix),
        answer=tf.gather(q["answer"], ix),
    )

@seqio.map_over_dataset
def scifi_explanation_only(ex):
    return dict(
        image=ex["image"],
        question=ex["question"],
        answer=ex["explanation"],
    )


def filter_named_entity(ds):
    @seqio.map_over_dataset
    def _load_image(ex):
        ex["image"] = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
        return ex

    ds = _load_image(ds)
    return ds.filter(lambda x: tf.reduce_min(tf.shape(x["image"])[:2]) >= 32)


@seqio.map_over_dataset()
def extract_named_entity(ex):
    qs = ex["questions"]
    return {
        "image": ex["image"],
        "metadata/image_url": ex["url"],
        "metadata/entity": ex["entity"],
        "questions": qs["question"],
        "answers": qs["answer"],
    }

@gin.configurable()
def extract_individual_vqa(ds, test=False, answer_mode="best"):

    @seqio.map_over_dataset(num_seeds=1)
    def _extract(ex, seed):
        if "questions" in ex:
            question = ex["questions"]
        else:
            question = ex["question"]
        out = dict(
            image=ex["image"],
            question=question,
        )
        out.update(_add_metadata(ex))
        out["metadata/question"] = question
        if ex.get("answers") is not None:
            out["metadata/references"] = tf.strings.reduce_join(ex["answers"], separator="\n")
        elif ex.get("answer") is not None:
            out["metadata/references"] = ex["answer"]

        if not test:
            if "answer" in ex:
                answer = ex["answer"]
            else:
                answer = ex["answers"]
            if answer.dtype in [tf.int32, tf.int64]:
                answer = tf.strings.as_string(answer)
            if len(answer.shape) == 1 and tf.shape(answer)[0] == 0:
                answer = tf.expand_dims("", 0)
            if len(answer.shape) == len(question.shape):
                pass
            # Handle questions with multiple answers
            elif answer_mode == "random":
                assert len(answer.shape) == 1
                answer = answer[tf.random.stateless_uniform((), seed, 0, tf.shape(answer)[0], dtype=tf.int32)]
            elif answer_mode == "best":
                def _get_best(_answer):
                    vals, _, counts = tf.unique_with_counts(_answer)
                    count_thresh = tf.reduce_max(counts)
                    vals = tf.boolean_mask(vals, counts >= count_thresh)
                    return vals[tf.random.stateless_uniform((), seed, 0, tf.shape(vals)[0], dtype=tf.int32)]
                if len(answer.shape) == 1:
                    answer = _get_best(answer)
                elif isinstance(answer, tf.RaggedTensor):
                    n = tf.shape(answer)[0]
                    answer_arr = tf.TensorArray(dtype=tf.string, size=n, element_shape=())
                    for i in range(n):
                        answer_arr = answer_arr.write(i, _get_best(answer[i]))
                    answer = answer_arr.stack()
                else:
                    answer = tf.map_fn(_get_best, answer)
            elif answer_mode == "all_segments":
                out["text"] = answer
            elif answer_mode == "all_segments_weighted":
                out["text"] = answer
                out["text_weights"] = 1.0 / tf.cast(tf.shape(answer)[-1], tf.float32)
            elif answer_mode == "all":
                if len(answer.shape) == 1:
                    answer = stateless_shuffle(answer, seed)
                    answer = tf.strings.reduce_join(answer, separator="\n", axis=-1)
                elif isinstance(answer, tf.RaggedTensor):
                    n = tf.shape(answer)[0]
                    answer_arr = tf.TensorArray(dtype=tf.string, size=n, element_shape=())
                    for i in range(n):
                        answer_arr = answer_arr.write(i, tf.strings.reduce_join(tf.random.shuffle(answer[i]), separator="\n", axis=-1))
                    answer = answer_arr.stack()
                else:
                    answer = tf.map_fn(tf.random.shuffle, answer)
                    answer = tf.strings.reduce_join(answer, separator="\n", axis=-1)
            else:
                raise NotImplementedError()
            out["text"] = answer
        return out
    return _extract(ds)


@seqio.map_over_dataset()
def extract_khan_academy(ex):
    return dict(
        image=ex["image"],
        image_url=ex["image_url"],
        prompt="Answer this question",
        text=ex["gptResponse"]
    )

@seqio.map_over_dataset()
def extract_vaia_qa_latex_image(ex, add_short_answer=False, set_short_answer_first=False):
    if ex["has_image"]:
        image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
        image = tf.expand_dims(image, 0)[:1]
    else:
        # image = get_blank_image() # blank image
        image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
        image = tf.expand_dims(image, 0)[:0]
    img_h = tf.shape(image)[1]
    img_w = tf.shape(image)[2]

    if add_short_answer:
        if set_short_answer_first:
            answer = tf.strings.join(["Answer: ", ex["short_answer"], "\n\n", ex["answer"]])
        else:
            answer = tf.strings.join([ex["answer"], "\n\n", "Answer: ", ex["short_answer"]])
    else:
        answer = ex["answer"]
    out = dict(
        image=image, # 4-d tensor
        text=answer,
        prompt=tf.strings.join([ex["latex_question"], "\n"]),
    )
    out["metadata/images"] = image
    out.update(_add_metadata(ex))
    out["metadata/batch_id"] = ex["batch_id"]
    out["metadata/image_size"] = [img_w, img_h]
    return out

@seqio.map_over_dataset()
def extract_vqa_online(ex):
    out = dict(
        image=ex["image"],
        prompt=tf.strings.join([ex["question"], "\n"]),
        text=ex["answer"]
    )
    out.update(_add_metadata(ex))
    out["metadata/row_id"] = ex["row_id"]
    return out


@seqio.map_over_dataset()
def extract_scifi_joint(ex):
    if "qa_pairs" in ex:
        q = ex["qa_pairs"]
    else:
        q = ex["qa"]
    prompts = tf.concat([["Describe this image in detail."], q["question"]], 0)
    responses = tf.concat([ex["summary"][None], q["answer"]], 0)
    return dict(
        image=ex["image"],
        prompt=prompts,
        text=responses,
    )


def remove_no_qa(ds):
    def _filter(ex):
        if "qa_pairs" in ex:
            q = ex["qa_pairs"]
        else:
            q = ex["qa"]
        return tf.shape(q["question"])[0] > 0
    return ds.filter(_filter)


@seqio.map_over_dataset()
def extract_scifi_qa_exp(ex):
    return dict(
        image=ex["image"],
        question=ex["question"],  # Array of questions
        answer=tf.strings.join([ex["explanation"], " Answer: ", ex["answer"]]),
    )


@seqio.map_over_dataset(num_seeds=1)
def extract_scifi_qa_demo(ex, seed):
    # if tf.random.stateless_uniform((), 0, 1) > 0.5:
    answer = tf.strings.join([ex["explanation"], " Answer: ", ex["answer"]])
    # else:
    #     answer = ex["explanation"]
    return dict(
        image=ex["image"],
        question=ex["question"],  # Array of questions
        answer=answer,
    )


@seqio.map_over_dataset()
def clock_bench_preprocessor(ex):
    out = dict(
        image=ex["image"],
        prompt="What time is being shown?",
    )
    for k in ["hour", "minute", "second", "answerable"]:
        out[f"metadata/{k}"] = ex[k]
    return out


def deg2rad(x):
    return x*math.pi/180.0


def get_affine_matrix(center, angle, translate, scale, shear):
    # From https://github.com/pytorch/vision/blob/f96c42fca53230057b16941b078a0a9eee06e20f/torchvision/transforms/functional.py#L1006
    rot = deg2rad(angle)
    sx = deg2rad(shear[0])
    sy = deg2rad(shear[1])

    cx, cy = center
    tx, ty = translate

    # RSS without scaling
    a = tf.cos(rot - sy) / tf.cos(sy)
    b = -tf.cos(rot - sy) * tf.tan(sx) / tf.cos(sy) - tf.sin(rot)
    c = tf.sin(rot - sy) / tf.cos(sy)
    d = -tf.sin(rot - sy) * tf.tan(sx) / tf.cos(sy) + tf.cos(rot)

    matrix = [a, b, 0.0, c, d, 0.0]
    matrix = [x * scale for x in matrix]
    # Apply inverse of center translation: RSS * C^-1
    matrix[2] += matrix[0] * (-cx) + matrix[1] * (-cy)
    matrix[5] += matrix[3] * (-cx) + matrix[4] * (-cy)
    # Apply translation and center : T * C * RSS * C^-1
    matrix[2] += cx + tx
    matrix[5] += cy + ty
    return matrix


def quantize_point(coor, max_dim, mode="percent-precision-1"):
    max_dim = tf.cast(max_dim, tf.float32)
    coor = tf.cast(coor, tf.float32)
    x = (coor / max_dim)
    if mode == "percent-precision-1":
        return tf.strings.as_string(x*100, precision=1)
    elif mode == "zero_to_one":
        return tf.strings.as_string(x, precision=3)
    elif mode == "1k":
        return tf.strings.as_string(x*1000, precision=0)
    else:
        raise NotImplementedError(mode)


def construct_pointing_format(label_text, alt_text, x_str, y_str):
    if alt_text is None:
        alt_text = label_text
    np = tf.shape(x_str)[0]
    if np == 0:
        output = ""
    elif np == 1:
        output = tf.strings.join([
            '<point x="', x_str[0], '" y="', y_str[0], '" alt="',
            alt_text, '">', label_text, '</point>'
        ])
    else:
        ids = tf.strings.as_string(tf.range(1, np + 1, dtype=tf.int32))
        xs = tf.strings.join(["x", ids, '="', x_str, '"'])
        ys = tf.strings.join(["y", ids, '="', y_str, '"'])
        points = tf.strings.reduce_join(tf.reshape(tf.stack([xs, ys], 1), [-1]), separator=' ', axis=-1)
        output = tf.strings.join(
            ["<points ", points, ' alt="', alt_text, '">', label_text, "</points>"])
    return output


def order_points(x, y, seed, point_order):
    if point_order == "natural":
        return x, y

    if point_order == "random":
        ix = stateless_permutation(tf.shape(x)[0], seed)
    elif point_order == "xy":
        x_float, y_float = tf.strings.to_number(x), tf.strings.to_number(y)
        ix = tf.argsort(x_float*100000 + y_float)
    elif point_order == "yx":
        x_float, y_float = tf.strings.to_number(x), tf.strings.to_number(y)
        ix = tf.argsort(y_float*100000 + x_float)
    else:
        raise NotImplementedError(point_order)
    return tf.gather(x, ix), tf.gather(y, ix)


@gin.configurable()
def points_to_text(x, y, w, h, seed, label=None, alt_text=None, point_mode="percent-precision-1",
                   point_order="xy", point_list_mode="tag"):
    """Returns a string encoding of a list of points"""
    x = quantize_point(x, w, point_mode)
    y = quantize_point(y, h, point_mode)
    # Order the quantized points to make the order matches what was generated, this can matter
    # when points have the same quantized value e.g, (10.001, 20) (10.002, 10) should be
    # represented (10, 10), (10, 20), but if we sort before quantization we get (10, 20), (10, 10)
    x, y = order_points(x, y, seed, point_order)
    if point_list_mode == "tag":
        return construct_pointing_format(label, alt_text, x, y)
    elif point_list_mode == "paren":
        n = tf.shape(x)[0]
        return tf.strings.reduce_join(tf.strings.join([
            "(", x, ", ", y, ")"
        ]), separator=", ")
        # if n == 0:
        #     output = ""
        # else:
        #     ids = tf.strings.as_string(tf.range(1, np + 1, dtype=tf.int32))
        #     xs = tf.strings.join(["x", ids, '="', x_str, '"'])
        #     ys = tf.strings.join(["y", ids, '="', y_str, '"'])
        #     points = tf.strings.reduce_join(tf.reshape(tf.stack([xs, ys], 1), [-1]), separator=' ', axis=-1)
        #     output = tf.strings.join(
        #         ["<points ", points, ' alt="', alt_text, '">', label_text, "</points>"])
        # return output
    else:
        raise NotImplementedError(point_list_mode)


def points_to_answer(x, y, w, h, seed, label, is_counting, alt_text=None):
    count = tf.shape(x)[0]
    if is_counting:
        if count == 0:
            return "There are none."
        else:
            point_text = points_to_text(x, y, w, h, seed, label, alt_text)
            return tf.strings.join([
                "Counting the ", point_text,
                " shows a total of ",
                tf.strings.as_string(count),
                "."
            ])
    else:
        if count == 0:
            return "There are none."
        else:
            return points_to_text(x, y, w, h, seed, label, alt_text)


@seqio.map_over_dataset(num_seeds=2)
def extract_point_qa(ex, seeds, answer_type="y_major"):
    ex["image"] = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]

    questions = ex["questions"]
    question = questions["question"]
    n = tf.shape(question)[0]
    answers = tf.TensorArray(tf.string, size=n, element_shape=())
    point_text = questions["annotations"]["point_text"]
    point_seeds = tf.RaggedTensor.from_row_splits(
        row_splits=point_text.row_splits,
        values=tf.random.split(seeds[0], num=tf.shape(point_text.values)[0])
    )
    for question_ix in range(n):
        anno = questions["annotations"]
        answer = questions["answer_with_placeholders"][question_ix]
        n_anno = tf.shape(anno["point_text"][question_ix])[0]
        for anno_ix in range(n_anno):
            points = anno["points"][question_ix, anno_ix]
            point_text = points_to_answer(
                points[:, 0], points[:, 1], 100, 100,
                point_seeds[question_ix, anno_ix],
                anno["point_text"][question_ix, anno_ix],
                False,
                alt_text=anno["alt_text"][question_ix, anno_ix],
            )
            answer_split = tf.strings.split(answer, sep="<|POINT|>", maxsplit=1)
            answer = tf.strings.join([answer_split[0], point_text, answer_split[1]])
        # Make sure all placeholders where used
        tf.debugging.assert_equal(tf.shape(tf.strings.split(answer, sep="<|POINT|>"))[0], 1)
        answers = answers.write(question_ix, answer)

    messages = tf.stack([question, answers.stack()], axis=1)
    messages = tf.reshape(messages, [-1])
    conversation_ids = tf.range(tf.shape(messages)[0] // 2, dtype=tf.int32)
    conversation_ids = tf.repeat(conversation_ids, 2)
    out = dict(
        image=ex["image"],
        messages=tf.RaggedTensor.from_value_rowids(messages, conversation_ids)
    )
    ix = stateless_permutation(tf.shape(messages)[0], seeds[1])
    messages = tf.gather(messages, ix)
    out.update(_add_metadata(ex))
    out["metadata/image_size"] = [img_w, img_h]
    return out


def select_point(mask):
    bs = tf.shape(mask)[0]
    valid = tf.cast(mask, tf.float32)
    h, w = tf.shape(mask)[1], tf.shape(mask)[2]
    ys = tf.range(h, dtype=tf.int32)
    xs = tf.range(w, dtype=tf.int32)

    n = tf.reduce_sum(valid, [1, 2])
    cy = tf.reduce_sum(tf.cast(ys[None, :, None], tf.float32) * valid, [1, 2]) / n  # [bs]
    cx = tf.reduce_sum(tf.cast(xs[None, None, :], tf.float32) * valid, [1, 2]) / n  # [bs]

    dist_y = tf.square(tf.range(h, dtype=tf.float32)[None, :] - cy[:, None])  # [bs, h]
    dist_x = tf.square(tf.range(w, dtype=tf.float32)[None, :] - cx[:, None])  # [bs, w]
    dist = dist_y[:, :, None] + dist_x[:, None, :]  # [batch, h, w]
    dist = dist + (1 - valid) * 1e12
    min_dist = tf.argmin(tf.reshape(dist, [bs, -1]), axis=-1)  # [batch]
    w = tf.cast(w, min_dist.dtype)
    cy = tf.cast(min_dist // w, tf.float32)
    cx = tf.cast(min_dist % w, tf.float32)
    return cx, cy


@seqio.map_over_dataset
def refexp_pointing(ex):
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]
    objects = ex["objects"]

    # Shuffle objects so what object gets truncated if the sequence gets truncated is randomized
    refexps = objects['refexp']['raw']
    bbox = objects["bbox"]
    mask = tf.squeeze(objects["mask"], -1)

    ix = tf.range(0, tf.shape(refexps)[0], dtype=tf.int32)
    ix = tf.random.shuffle(ix)
    refexps = tf.gather(refexps, ix)
    bbox = tf.gather(bbox, ix)
    mask = tf.gather(mask, ix)

    cx, cy = select_point(mask)
    answers = points_to_text(img_h, img_w, cx, cy)

    out = {
        "image": ex["image"],
        "refexp": refexps.values,
        "metadata/image_size": tf.stack([img_w, img_h,]),
        "text": tf.repeat(answers, refexps.row_lengths()),
    }
    if "image_url" in ex:
        out["metadata/image_url"] = ex["image_url"]
    return out


@seqio.map_over_dataset
def refexp_pointing_inf(ex):
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]

    objects = ex["objects"]
    mask = tf.squeeze(objects["mask"], -1)
    cx, cy = select_point(mask)
    answers = points_to_text(img_h, img_w, cx, cy)

    refexps = objects["refexp"]["raw"]

    # We can't use `mask` directly since it is variable size, and thus it
    # will break batching. Here we serialize it instead
    serialized_masks = tf.map_fn(tf.io.serialize_tensor, mask, fn_output_signature=tf.string)
    out = {
        "image": ex["image"],
        "refexp": refexps,
        "metadata/bbox": objects["bbox"],
        "metadata/answer": answers,
        "metadata/mask": serialized_masks,
        "metadata/image_size": tf.stack([img_w, img_h]),
    }
    out.update({k: v for k, v in ex.items() if k.startswith("metadata/")})
    return out

@seqio.map_over_dataset
def extract_andriod_control_inf(ex, mode):
    if mode == "ll":
        prompt = tf.strings.join(["low_level: ", ex["metadata/ll_instruction"]])
    elif mode == "hl_ll":
        prompt = tf.strings.join([
            "high_level: ", ex["metadata/hl_instruction"],
            " low_level: ", ex["metadata/ll_instruction"]
        ])
    elif mode == "hl":
        prompt = tf.strings.join(["high_level: ", ex["metadata/hl_instruction"]])
    elif mode == "hl_cot":
        prompt = tf.strings.join(["high_level_cot: ", ex["metadata/hl_instruction"]])
    else:
        raise NotImplementedError()

    out = dict(
        image=ex["image"],
        prompt=prompt,
        text=ex["metadata/target_action"]
    )
    out.update(_add_metadata(ex))
    return out

@seqio.map_over_dataset
def extract_android_control(ex):
    # Each image has three tasks:
    # low level -> action
    # high+low level -> action
    # high level -> action
    # high level -> low level + action (CoT)
    out = dict(
        image=ex["image"],
        prompt=tf.stack([
            tf.strings.join(["low_level: ", ex["metadata/ll_instruction"]]),
            tf.strings.join([
                "high_level: ", ex["metadata/hl_instruction"],
                " low_level: ", ex["metadata/ll_instruction"]
            ]),
            tf.strings.join(["high_level: ", ex["metadata/hl_instruction"]]),
            tf.strings.join(["high_level_cot: ", ex["metadata/hl_instruction"]]),
        ]),
        text=tf.stack([
            ex["metadata/target_action"],
            ex["metadata/target_action"],
            ex["metadata/target_action"],
            tf.strings.join(["Plan: ", ex["metadata/ll_instruction"], " Action: ", ex["metadata/target_action"]]),
        ])
    )
    # Only needed if visualizing
    # ex["image"] = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    # img_h = tf.shape(ex["image"])[0]
    # img_w = tf.shape(ex["image"])[1]
    # out["metadata/image_size"] = tf.stack([img_w, img_h,])
    out.update(_add_metadata(ex))
    return out


@seqio.map_over_dataset(num_seeds=1)
def refexp(ex, seed):
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]
    objects = ex["objects"]

    # Shuffle objects so what object gets truncated if the sequence gets truncated is randomized
    refexps = objects['refexp']['raw']
    bbox = objects["bbox"]
    ix = stateless_permutation(tf.shape(refexps)[0], seed)
    refexps = tf.gather(refexps, ix)
    bbox = tf.gather(bbox, ix)

    x2 = bbox[:, 0] + bbox[:, 2]
    y2 = bbox[:, 1] + bbox[:, 3]
    with tf.control_dependencies([
        tf.debugging.assert_equal(tf.reduce_any(x2 <= tf.cast(img_w, tf.float32)), True),
        tf.debugging.assert_equal(tf.reduce_any(y2 <= tf.cast(img_h, tf.float32)), True)
    ]):
        answers = points_to_text(
            img_h, img_w,
            tf.reshape(tf.stack([bbox[:, 0], x2], 1), [-1]),
            tf.reshape(tf.stack([bbox[:, 1], y2], 1), [-1]))
        answers = tf.strings.reduce_join(tf.reshape(answers, [-1, 2]), separator=" ", axis=1)

    out = {
        "image": ex["image"],
        "refexp": refexps.values,
        "metadata/bbox": bbox,
        "metadata/image_size": tf.stack([img_w, img_h,]),
        "text": tf.repeat(answers, refexps.row_lengths()),
    }

    if "image_url" in ex:
        out["image_url"] = ex["image_url"]
    return out


@seqio.map_over_dataset
def refexp_inf(ex):
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]
    out = {
        "image": ex["image"],
        "refexp": ex["objects"]["refexp"]["raw"],
        "metadata/bbox": ex["objects"]["bbox"],
        "metadata/image_size": tf.stack([img_w, img_h,]),
    }
    out.update({k: v for k, v in ex.items() if k.startswith("metadata/")})
    return out


def point_text_interleaved(*args):
    raise NotImplementedError()


@seqio.map_over_dataset
def web_pointing_preprocessor(ex):
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]

    question = point_text_interleaved(
        img_h, img_w, ex["question"], ex["question_points"]["x"], ex["question_points"]["y"])
    answer = point_text_interleaved(
        img_h, img_w, ex["answer"], ex["answer_points"]["x"], ex["answer_points"]["y"])
    answer_points = tf.stack([ex["answer_points"]["x"], ex["answer_points"]["y"]], axis=1)
    return {
        "question": question,
        "answer": answer,
        "image": ex["image"],
        "metadata/image_size": [img_w, img_h],
        "metadata/question_type": ex["question_type"],
        "metadata/answer_points": tf.io.serialize_tensor(answer_points),
        "metadata/answer": answer,
    }


def filter_pointing(ds):
    return ds.filter(lambda ex: tf.shape(ex["answer_points"]["x"])[0] >= 1)


def filter_qa(ds):
    return ds.filter(lambda ex: tf.shape(ex["answer_points"]["x"])[0] == 0)

# vaia filtering
def filter_image_only(ds):
    return ds.filter(lambda ex: ex["has_image"])

def filter_mc(ds):
    return ds.filter(lambda ex: ex["is_mc"])

def remove_is_long(ds):
    return ds.filter(lambda ex: not ex["is_long"])

def remove_has_multiple_parts(ds):
    return ds.filter(lambda ex: not ex["has_multiple_parts"])


def _split(ds: tf.data.Dataset, keys, n_splits=2):
    def _map(ex):
        n = tf.shape(ex[keys[0]])[0]
        if n < n_splits:
            return tf.data.Dataset.from_tensors(ex)
        else:
            # import pdb; pdb.set_trace()
            bs = n // n_splits
            remainder = n - bs*n_splits
            lens = tf.concat([
                tf.ones([remainder], dtype=tf.int32),
                tf.zeros([n_splits-remainder], dtype=tf.int32),
            ], axis=0) + bs
            tf.debugging.assert_equal(tf.reduce_sum(lens), n)
            ends = tf.cumsum(lens)

            parts = []
            for split_ix in range(n_splits):
                part_ex = dict(ex)
                e = ends[split_ix]
                s = e - lens[split_ix]
                for k in keys:
                    if isinstance(k, tuple):
                        assert len(k) == 2
                        part_ex[k[0]][k[1]] = ex[k[0]][k[1]][s:e]
                    else:
                        part_ex[k] = ex[k][s:e]
                parts.append(part_ex)

            ds = tf.data.Dataset.from_tensors(parts[0])
            for sub_ds in parts[1:]:
                sub_ds = tf.data.Dataset.from_tensors(sub_ds)
                ds = ds.concatenate(sub_ds)
            return ds

    return ds.flat_map(_map)



def split(ds, n=2):
    # return ds
    return _split(ds, [k for k in [
        "question",
        "label",
        "text",
        "entity",
        "messages"
    ] if k in ds.element_spec], n_splits=n)


def split_points(ds, max_points=50):
    label = "question" if "question" in ds.element_spec else "label"
    return _split(ds, [
        "question", label, "notInImage",
        ("answer_points", "x"),
        ("answer_points", "y"),
    ])


@seqio.map_over_dataset
def fix_count_qa(ex):
    ex["label"] = ex["label"][::2]
    tf.debugging.assert_equal(tf.shape(ex["answer_points"]["x"])[0], tf.shape(ex["label"])[0])
    return ex


def filter_points(ds, max_number=40):

    def _add_valid(ex):
        valid = (
            tf.reduce_all(ex["answer_points"]["x"] >= 0.0, axis=-1) &
            tf.reduce_all(ex["answer_points"]["x"] <= 100.0, axis=-1) &
            tf.reduce_all(ex["answer_points"]["y"] >= 0.0, axis=-1) &
            tf.reduce_all(ex["answer_points"]["y"] <= 100.0, axis=-1) &
            (ex["answer_points"]["y"].row_lengths() <= max_number)
        )
        ex["valid"] = valid
        return ex
    ds = ds.map(_add_valid)
    ds = ds.filter(lambda ex: tf.reduce_any(ex["valid"]))
    return ds


# def filter_points(ds, max_number=30):
# n_points = ds["answer_points"]["x"].row_lengths()
# parts = tf.TensorArray(tf.int32, size=tf.shape(n_points[0]), element_shape=tf.TensorShape([None]))
# total = 0
# on_row = 0
# for i in range(n_points):
#     n = n_points[i]
#     if n > max_number:
#         continue
#     if n + total > max_number:
#
# return ds


@seqio.map_over_dataset(num_seeds=2)
def pointing_preprocessor(ex, sequence_length, seeds, with_count=False):
    image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    img_h = tf.shape(image)[0]
    img_w = tf.shape(image)[1]

    ix = tf.where(ex["valid"])[:, 0]
    ix = stateless_shuffle(ix, seeds[0])
    if "label" in ex:
        question = tf.strings.lower(ex["label"])
    else:
        question = ex["question"]
    question = tf.gather(question, ix)  # [n_question]
    points_x = tf.gather(ex["answer_points"]["x"], ix)  # [n_question, n_points[ragged]]]
    points_y = tf.gather(ex["answer_points"]["y"], ix)
    not_in_image = tf.gather(ex["notInImage"], ix)  # [n_question]

    n = tf.shape(points_x)[0]
    point_text = tf.TensorArray(dtype=tf.string, size=n, element_shape=())  # [n_question]
    point_seeds = tf.random.split(seeds[1], n)
    for i in range(n):
        answer = points_to_answer(points_x[i], points_y[i], 100, 100, point_seeds[i], question[i], with_count)
        point_text = point_text.write(i, answer)
    return {
        "image": image,
        "metadata/image_size": [img_w, img_h],
        "entity": question,
        "question": question,
        "text": point_text.stack(),
    }


@seqio.map_over_dataset
def pointing_inf_preprocessor(ex):
    ex["image"] = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    img_h = tf.shape(ex["image"])[0]
    img_w = tf.shape(ex["image"])[1]

    question = ex["question"]
    not_in_image = tf.shape(ex["answer_points"]["x"])[0] == 0

    # points are stored in normalized format, de-normalize here
    points_x = ex["answer_points"]["x"] * tf.cast(img_w, tf.float32) / 100.0
    points_y = ex["answer_points"]["y"] * tf.cast(img_h, tf.float32) / 100.0

    out = dict(
        image=ex["image"],
        question=question,
        entity=question,
    )
    out.update(_add_metadata(ex))
    out["metadata/not_in_image"] = not_in_image
    # We can't use `mask` directly since it is variable size, and thus it
    # will break batching. Here we serialize it instead
    serialized_masks = tf.map_fn(tf.io.serialize_tensor, ex["masks"], fn_output_signature=tf.string)
    serialized_masks = tf.strings.reduce_join(serialized_masks, separator="|||")
    out["metadata/mask"] = serialized_masks
    out["metadata/question"] = question
    out["metadata/answer_points"] = tf.io.serialize_tensor(tf.stack([points_x, points_y], 1))
    out["metadata/image_size"] = [img_w, img_h]

    return out


@seqio.map_over_dataset(num_seeds=1)
def count_qa_preprocessor_inf(ex, sequence_length, seed):
    image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    img_h = tf.shape(image)[0]
    img_w = tf.shape(image)[1]

    entity = tf.strings.substr(
        ex["question"], len("How many "), tf.strings.length(ex["question"]) - len("How many "))
    entity = tf.strings.split(entity, sep=" are ", maxsplit=1)[0]
    entity = tf.strings.lower(entity)
    tf.debugging.assert_equal(tf.strings.length(entity) != 0, True)

    return {
        "image": image,
        "metadata/image_size": [img_w, img_h],
        "metadata/count": tf.strings.to_number(ex["answer"]),
        "question": ex["question"],
        "entity": entity,
    }


@seqio.map_over_dataset(num_seeds=1)
def count_qa_preprocessor(ex, sequence_length, seed, with_count=False,
                          for_inference=False):
    point_answer = ex["point_answer"]
    numbers_str = tf.strings.regex_replace(point_answer, r'\.$', '')
    numbers_str = tf.strings.regex_replace(numbers_str, r'[^\d\.\s]+', '')
    numbers_str = tf.strings.strip(numbers_str)
    numbers = tf.strings.split(numbers_str)
    float_numbers = tf.strings.to_number(numbers, out_type=tf.float32)
    coordinates = tf.reshape(float_numbers, (-1, 3))
    points_x = coordinates[:, 1]
    points_y = coordinates[:, 2]

    image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    img_h = tf.shape(image)[0]
    img_w = tf.shape(image)[1]
    entity = tf.strings.substr(
        ex["question"], len("How many "), tf.strings.length(ex["question"]) - len("How many "))
    entity = tf.strings.split(entity, sep=" are ", maxsplit=1)[0]
    entity = tf.strings.lower(entity)
    tf.debugging.assert_equal(tf.strings.length(entity) != 0, True)
    count = tf.strings.to_number(ex["answer"], out_type=tf.int32)
    if for_inference:
        return {
            "image": image,
            "metadata/image_size": [img_w, img_h],
            "metadata/count": count,
            "question": ex["question"],
            "entity": entity,
        }
    else:
        tf.debugging.assert_equal(count, tf.shape(points_x)[0])
        # points are already normalized so use w=1, h=1
        answer = points_to_answer(points_x, points_y, 1, 1, seed, entity, with_count)
        return {
            "image": image,
            "metadata/image_size": [img_w, img_h],
            "metadata/count": count,
            "question": ex["question"],
            "entity": entity,
            "text": answer,
        }


@gin.configurable()
@seqio.map_over_dataset
def cleanup_preprocessor(ex, preprocess=False):
    if preprocess:
        ex["prompt"] = tf.strings.join(
            [
                "[[User]]: Correct the spelling and punctuation mistakes on the following transcript based on what appears in the image.\n\n{before} ",
                ex["prompt"],
                "\n[[Assistant]]: {after}"
            ]
        )
        return ex
    else:
        return ex


@gin.configurable()
@seqio.map_over_dataset
def random_text_preprocessor(ex, preprocess=False):
    ex["prompt"] = "What does the text say in this image?"
    if preprocess:
        ex["prompt"] = tf.strings.join(["[[User]]: ", ex["prompt"], "\n[[Assistant]]:"])
        return ex
    else:
        return ex


@seqio.map_over_dataset(num_seeds=25)
def clock_augmentation(ex, seeds):
    seeds = list(seeds)
    image = ex["image"]

    # Apply shear, rotation, and scale through one affine matrix
    height = tf.cast(tf.shape(image)[0], tf.float32)
    width = tf.cast(tf.shape(image)[1], tf.float32)

    _call_id = [0]

    def _rng(_minval=0, _maxval=1, shape=(), dtype=tf.float32):
        return tf.random.stateless_uniform(shape, seeds.pop(), _minval, _maxval, dtype=dtype)

    sel = _rng(0, 1)
    if sel < 0.1:
        # Straight on
        shear_x = 0.
        shear_y = 0.
        rotation = 0.
    elif sel < 0.5:
        # Normal looking
        shear_x = _rng(-10, 10)
        shear_y = _rng(-10, 10)
        rotation = _rng(-25, 25)
    else:
        # Allowed to be very wonky
        # if tf.random.stateless_uniform((), seeds.pop(), 0, 1) > 0.8:
        #     image = image[:, ::-1]

        if _rng() > 0.5:
            shear_x = _rng( -30, 30)
            shear_y = _rng( -30, 30)
        else:
            shear_x = _rng( -10, 10)
            shear_y = _rng( -10, 10)
        rng = _rng( 0, 1)
        if rng < 0.2:
            rotation = _rng( -25, 25)
        elif rng < 0.6:
            rotation = _rng( -80, 80)
        else:
            rotation = _rng( -180, 180)

    if _rng() > 0.5:
        scale = _rng( 0.3, 2)
    else:
        scale = _rng( 0.3, 1)
    # Pad so upscaling/rotation will not move the image out of bounds
    pad = tf.cast(tf.maximum(height, width)*0.5, tf.int32)
    image = tf.pad(image, [[pad, pad], [pad, pad], [0, 0]], constant_values=1)
    height = tf.cast(tf.shape(image)[0], tf.float32)
    width = tf.cast(tf.shape(image)[1], tf.float32)

    image = tf.keras.ops.image.affine_transform(
        image,
        tf.stack(get_affine_matrix(
            [height/2, width/2],
            rotation,
            [0, 0],
            1/scale,
            [shear_x, shear_y]
        ) + [0., 0.]),
        interpolation='bilinear',
        fill_mode='constant',
        fill_value=1.,
        data_format='channels_last'
    )

    # Crop, otherwise it would be impossible to put the image at the corner of the image
    not_white = tf.logical_not(tf.reduce_all(image > 0.99, -1))
    no_white_ix = tf.where(not_white)
    top_left = tf.reduce_min(no_white_ix, axis=0)
    bottom_right = tf.reduce_max(no_white_ix, axis=0)
    image = tf.image.crop_to_bounding_box(
        image,
        offset_height=tf.cast(top_left[0], tf.int32),
        offset_width=tf.cast(top_left[1], tf.int32),
        target_height=tf.cast(bottom_right[0] - top_left[0] + 1, tf.int32),
        target_width=tf.cast(bottom_right[1] - top_left[1] + 1, tf.int32),
    )

    # Translate
    height, width = tf.shape(image)[0], tf.shape(image)[1]
    translation_seed = _rng(0, 1)
    if translation_seed < 0.2:
        h_pad = _rng(0, height//2, (2,), dtype=tf.int32)
        w_pad = _rng(0, width//2, (2,), dtype=tf.int32)
    else:
        h_pad = _rng(0, height*2, (2,), dtype=tf.int32)
        w_pad = _rng(0, width*2, (2,), dtype=tf.int32)
    image = tf.pad(image, [[h_pad[0], w_pad[0]], [h_pad[1], w_pad[1]], [0, 0]],
                   constant_values=1)

    # Random background color
    # color_rng = tf.random.stateless_uniform((4,), seeds.pop(), 0, 1)
    # random_color = color_rng[:3]
    # valid = tf.reduce_all(tf.reduce_sum(tf.abs(random_color[None, None, :] - image), -1) > 0.03)
    # if color_rng[0] < 0.2 and valid:
    #     image = tf.where(tf.reduce_all(image < 0.99, axis=-1, keepdims=True),
    #                      image, image * 0 + random_color[None, None, :])

    # Mild color hitter
    image = tf.image.stateless_random_hue(image, max_delta=0.05, seed=seeds.pop())
    image = tf.image.stateless_random_brightness(image, max_delta=0.15, seed=seeds.pop())
    image = tf.image.stateless_random_saturation(image, 0.8, 1.2, seed=seeds.pop())
    image = tf.image.stateless_random_contrast(image, 0.8, 1.2, seed=seeds.pop())

    # ex["metadata/unaugmented_image"] = ex["image"]
    ex["image"] = image
    return ex


@seqio.map_over_dataset
def clocks_preprocessor(ex):
    time_format = ex["time_format"]
    shows_seconds = ex["shows_seconds"]
    hour, minute, second = [tf.cast(ex[k], tf.int32) for k in ["hour", "minute", "second"]]
    if hour == 0:  # Midnight of the previous day
        am_pm = "PM"
        hour_str = 12
        hour = 24
    elif hour > 12:
        am_pm = "PM"
        hour_str = hour - 12
    else:
        hour_str = hour
        am_pm = "AM"
    hour_str = tf.strings.as_string(hour_str)
    minute_str = tf.strings.as_string(minute)
    if tf.strings.length(minute_str) == 1:
        minute_str = tf.strings.join(["0", minute_str])

    second_str = tf.strings.as_string(second)
    if tf.strings.length(second_str) == 1:
        second_str = tf.strings.join(["0", second_str])

    prefix = "The time shown is "

    if time_format == "The time is not shown":
        text = "The time is not shown in the image."
        hour, minute, second = -1, -1, -1
    else:
        if not shows_seconds:
            second = -1
        if time_format == "12 hour clock (without AM/PM)" and shows_seconds:
            if hour > 12:
                hour = hour - 12
            time = tf.strings.join([hour_str, ":", minute_str, ":", second_str])
        elif time_format == "12 hour clock (with AM/PM)" and shows_seconds:
            time = tf.strings.join([hour_str, ":", minute_str, ":", second_str, " ", am_pm])
        elif time_format == "12 hour clock (with AM/PM)" and not shows_seconds:
            time = tf.strings.join([hour_str, ":", minute_str, " ", am_pm])
        elif time_format == "12 hour clock (without AM/PM)" and not shows_seconds:
            if hour > 12:
                hour = hour - 12
            time = tf.strings.join([hour_str, ":", minute_str])
        else:
            time = ""  # Should never occur, but needed for tf analysis
        tf.debugging.assert_equal(tf.strings.length(time) > 0, True)
        text = tf.strings.join(["The time shown is ", time])
    image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    image = tf.image.convert_image_dtype(image, tf.float32)[:-120]  # remove the black shadow at the bottom
    return {
        "image": image,
        "prompt": "What time is being shown?",
        "text": text,
        "metadata/time_format": time_format,
        "metadata/hour": hour,
        "metadata/minute": minute,
        "metadata/text": text,
        "metadata/second": second,
    }


@seqio.map_over_dataset()
def atlas_obscura_preprocessor(ex):
    out = dict(
        image=ex["image"],
        prompt="Where was this picture taken?",
        text=tf.strings.join([
            ex["place"],
            " in ",
            ex["city"]
        ])
    )
    out["metadata/image_url"] = ex["image_url"]
    out["metadata/references"] = out["text"]
    return out


@seqio.map_over_dataset()
def famous_birthdays_preprocessor(ex):
    out = dict(
        image=ex["image"],
        image_url=ex["image_url"],
        prompt="Who is this?",
        text=ex["name"]
    )
    out["metadata/references"] = out["text"]
    return out


@seqio.map_over_dataset()
def mild_color_aug_preprocessor(ex):
    if "image_url" in ex:  # URL won't show the augmentations
        del ex["image_url"]
    # ex["metadata/unaugmented_image"] = ex["image"]
    ex["image"] = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    ex["image"] = mild_color_aug(ex["image"])
    return ex


def build_text_with_points(text, points, img_h, img_w):
    points = points_to_text(img_h, img_w, points[:, 0], points[:, 1])
    parts = tf.strings.split(text, sep="<ANS>")
    with_points = tf.strings.reduce_join(tf.reshape(tf.stack([
        parts,
        tf.pad(points, [[0, 1]], constant_values=""),
    ], 1), [-1]), separator="")
    return tf.strings.split(with_points, "\n\n")


@seqio.map_over_dataset()
def synth_count_preprocessor(example):
    image_shape = tf.shape(example["image"])
    h, w = image_shape[0], image_shape[1]
    questions = build_text_with_points(example["questions"], example["question_points"], h, w)
    answers = build_text_with_points(example["answers"], example["answer_points"], h, w)
    keep_q = tf.strings.regex_full_match(questions, "How many.*")
    keep_ans = tf.strings.regex_full_match(answers, "There are [0-9]+.*")
    keep = tf.logical_and(keep_q, keep_ans)
    questions = tf.boolean_mask(questions, keep)
    answers = tf.boolean_mask(answers, keep)
    ix = tf.range(0, tf.shape(answers)[0], dtype=tf.int32)
    ix = tf.random.shuffle(ix)
    return dict(
        image=example["image"],
        prompt=tf.gather(questions, ix),
        text=tf.gather(answers, ix),
    )


def synth_count_inf_preprocessor(ds):

    @seqio.map_over_dataset(num_seeds=1)
    def get_two(example, seed):
        image_shape = tf.shape(example["image"])
        h, w = image_shape[0], image_shape[1]
        questions = build_text_with_points(example["questions"], example["question_points"], h, w)
        answers = build_text_with_points(example["answers"], example["answer_points"], h, w)
        keep_q = tf.strings.regex_full_match(questions, "How many.*")
        keep_ans = tf.strings.regex_full_match(answers, "There are [0-9]+.*")
        keep = tf.logical_and(keep_q, keep_ans)
        questions = tf.boolean_mask(questions, keep)
        answers = tf.boolean_mask(answers, keep)

        ix = stateless_permutation(tf.shape(answers)[0], seed)[:2]
        return {
            "image": example["image"],
            "prompt": tf.gather(questions, ix),
            "metadata/references": tf.gather(answers, ix),
        }

    ds = get_two(ds)
    return flatten_parts(ds, ["prompt", "metadata/references"])


def mild_color_aug(image):
    image = tf.image.random_hue(image, max_delta=0.05)
    image = tf.image.random_brightness(image, max_delta=0.15)
    image = tf.image.random_saturation(image, 0.7, 1.3)
    image = tf.image.random_contrast(image, 0.8, 1.2)
    return image


@seqio.map_over_dataset()
def name_entity_augmentation(ex, p_high_color=0.7):
    ex["image"] = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
    image = ex["image"]
    image = tf.image.convert_image_dtype(image, tf.float32)

    # Horizontal flip
    if tf.random.uniform((), 0, 1) > 0.85:
        image = image[:, ::-1]

    # Random crop
    height = tf.cast(tf.shape(image)[0], tf.float32)
    width = tf.cast(tf.shape(image)[1], tf.float32)
    crop_rng = tf.random.uniform((), 0, 1)
    if crop_rng < 0.2:
        pass
    else:
        if crop_rng < 0.4:
            h_crop = height * 0.15
            w_crop = width * 0.15
        else:
            h_crop = height * 0.4
            w_crop = width * 0.4
        crop_h = tf.cast(tf.random.uniform((2,), 0, h_crop/2), tf.int32)
        crop_w = tf.cast(tf.random.uniform((2,), 0, w_crop/2), tf.int32)
        image = image[crop_h[0]:-crop_h[1]-1, crop_w[0]:-crop_w[1]-1]
        height = tf.cast(tf.shape(image)[0], tf.float32)
        width = tf.cast(tf.shape(image)[1], tf.float32)

    if tf.random.uniform(()) > p_high_color:
        image = tf.image.random_hue(image, max_delta=0.05)
        image = tf.image.random_brightness(image, max_delta=0.15)
        image = tf.image.random_saturation(image, 0.7, 1.3)
        image = tf.image.random_contrast(image, 0.8, 1.2)
    else:
        image = tf.image.random_hue(image, max_delta=0.1)
        image = tf.image.random_brightness(image, max_delta=0.3)
        image = tf.image.random_saturation(image, 0.0, 2.0)
        image = tf.image.random_contrast(image, 0.2, 1.5)

    # Apply shear, rotation, and scale through one affine matrix
    sel = tf.random.uniform((), 0, 1)
    if sel < 0.1:
        pass
    else:
        if sel < 0.15:  # Scale only
            shear_x = 0
            shear_y = 0
            rotation = 0
        if sel < 0.7:  # Mild
            shear_x = tf.random.uniform((), -2, 2)
            shear_y = tf.random.uniform((), -2, 2)
            rotation = tf.random.uniform((), -5, 5)
        else:  # Severe
            shear_x = tf.random.uniform((), -10, 10)
            shear_y = tf.random.uniform((), -10, 10)
            rotation = tf.random.uniform((), -20, 20)

        max_scale = 1.2
        scale = tf.random.uniform((), 0.4, max_scale)

        # Pad so upscaling/rotation will not move the image out of bounds
        pad = tf.cast(tf.maximum(height, width)*0.2, tf.int32)
        image = tf.pad(image, [[pad, pad], [pad, pad], [0, 0]], constant_values=1)

        image = tf.keras.ops.image.affine_transform(
            image,
            tf.stack(get_affine_matrix(
                [height/2, width/2],
                rotation,
                [0, 0],
                1/scale,
                [shear_x, shear_y]
            ) + [0., 0.]),
            interpolation='bilinear',
            fill_mode='constant',
            fill_value=1.,
            data_format='channels_last'
        )

    # Crop, otherwise it would be impossible to put the image at the corner of the image
    not_white = tf.logical_not(tf.reduce_all(image > 0.99, -1))
    no_white_ix = tf.where(not_white)
    top_left = tf.reduce_min(no_white_ix, axis=0)
    bottom_right = tf.reduce_max(no_white_ix, axis=0)

    # Very low chance center crop will get nothing but white space, we just skip
    if (
        (bottom_right[0] - top_left[0]) > 1 and (bottom_right[1] - top_left[1]) > 1
    ):
        image = tf.image.crop_to_bounding_box(
            image,
            offset_height=tf.cast(top_left[0], tf.int32),
            offset_width=tf.cast(top_left[1], tf.int32),
            target_height=tf.cast(bottom_right[0] - top_left[0] + 1, tf.int32),
            target_width=tf.cast(bottom_right[1] - top_left[1] + 1, tf.int32),
        )

    # Translate
    height, width = tf.shape(image)[0], tf.shape(image)[1]
    if tf.random.uniform((), 0, 1) < 0.1:
        h_pad = tf.zeros((2,), dtype=tf.int32)
        w_pad = tf.zeros((2,), dtype=tf.int32)
    elif tf.random.uniform((), 0, 1) < 0.8:
        h_pad = tf.random.uniform((2,), 0, 50, dtype=tf.int32)
        w_pad = tf.random.uniform((2,), 0, 50, dtype=tf.int32)
    else:
        pad = tf.cast(tf.maximum(height, width), tf.int32)
        h_pad = tf.random.uniform((2,), 0, pad, dtype=tf.int32)
        w_pad = tf.random.uniform((2,), 0, pad, dtype=tf.int32)
    image = tf.pad(image, [[h_pad[0], w_pad[0]], [h_pad[1], w_pad[1]], [0, 0]],
                   constant_values=1)

    if "image_url" in ex:  # URL won't show the augmentations
        del ex["image_url"]
    # ex["metadata/unaugmented_image"] = ex["image"]
    ex["image"] = image
    return ex


@seqio.map_over_dataset()
def wiki_art_preprocessor(ex):
    out = dict(
        image=ex["image"],
        prompt="What is this?",
        text=ex["question"]
    )
    out["metadata/title"] = ex["title"]
    out["metadata/gt"] = ex["question"]
    out["metadata/artist"] = ex["artist"]
    out["metadata/painting_url"] = ex["painting_url"]
    # if "metadata/unaugmented_image" in ex:
    #     out["metadata/unaugmented_image"] = ex["metadata/unaugmented_image"]
    return out

@seqio.map_over_dataset()
def oscar_preprocessor(ex):
    out = dict(
        image=ex["image"],
        prompt=ex["question"]
    )
    out.update(_add_metadata(ex))
    out["metadata/question"] = ex["question"]
    out["metadata/answer"] = ex["answer"]
    out["metadata/category"] = ex["category"]
    return out


@seqio.map_over_dataset()
def tulu_preprocessor(ex):
    return {
        "messages": ex["messages"]["content"],
    }
    # logging.info("Debugging tulue")
    # return {"messages": ex["messages"]["content"], "text_weights": 1e-6}


WIKI_DATA_QUESTION = "What is this? Respond with just a proper name."


@seqio.map_over_dataset()
def extract_wiki_data(ex):
    return dict(
        image=ex["image"],
        image_url=ex["image_url"],
        prompt=[
            WIKI_DATA_QUESTION,
            "What is this? Respond with the proper name of the main focus of the image and a few details about it."
        ],
        text=[
            tf.strings.strip(tf.strings.regex_replace(ex["question"], r"\(.*\)", "")),
            ex["gptResponse"],
        ]
    )


@seqio.map_over_dataset()
def extract_wiki_data_name(ex):
    target = tf.strings.strip(tf.strings.regex_replace(ex["question"], r"\(.*\)", ""))
    out = dict(
        image=ex["image"],
        image_url=ex["image_url"],
        prompt=WIKI_DATA_QUESTION,
        text=target,
    )
    out["metadata/references"] = target
    return out


@seqio.map_over_dataset()
def extract_wiki_data_describe(ex):
    out = dict(
        image=ex["image"],
        image_url=ex["image_url"],
        prompt="What is this? Respond with the proper name of the main focus of the image and a few details about it.",
    )
    out["metadata/references"] = ex["gptResponse"]
    return out


@gin.configurable()
def format_multiple_style_qa(ds, types=['multiple_choice', 'short_answer'], styles=['ai2_diagram', 'vqa2'], default_style='vqa2',
                             strip_instruction=False):
    def _extract(ex):
        prompt = ex["question"]
        out = dict(image=ex["image"])
        out.update(_add_metadata(ex))

        out["text"] = ex["answer"]
        out["metadata/references"] = ex["answer"]

        if ex["metadata/question_type"] == 'multiple_choice':
            style = styles[0]
        else:
            style = styles[1]
        if strip_instruction:
            if ex["metadata/question_type"] == "multiple_choice":
                # parts = tf.strings.split(prompt, "\n")
                # parts 1 is blank and part -1 is the instruction
                # prompt = tf.strings.reduce_join(tf.concat([parts[:1], parts[2:-1]], 0), separator="\n")
                prompt = prompt
            else:
                prompt = tf.strings.split(prompt, "\n")[0]

        out["style"] = style
        out["prompt"] = prompt
        return out
    ds = ds.map(_extract, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return ds


@gin.configurable()
def extract_mmmu(ds, types=['multiple-choice', 'open'], styles=['ai2_diagram', 'vqa2'], default_style='ai2_diagram', option_format="abc"):
    assert option_format == "abc"
    keys_tensor = tf.constant(types, dtype=tf.string)
    values_tensor = tf.constant(styles, dtype=tf.string)
    table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(keys_tensor, values_tensor),
        default_value=tf.constant(default_style, dtype=tf.string),
    )
    def _extract(ex):
        out = dict(image=tf.expand_dims(ex["image_1"], 0))
        out.update(_add_metadata(ex))
        style = table.lookup(ex["metadata/question_type"])
        out["style"] = style
        out["text"] = ex["answer"]
        out["metadata/references"] = ex["answer"]

        if style == styles[0]:
            abc = tf.constant(list("abcdefghi".upper()))
            options = ex["options"]
            num_options = tf.shape(options)[0]
            dummy_options = tf.tile(tf.constant([""], dtype=tf.string), [9 - num_options])
            out["metadata/options"] = tf.concat([options, dummy_options], axis=0)
            out["metadata/options"] = tf.ensure_shape(out["metadata/options"], [9])

            short_options = abc[:num_options]
            options = tf.stack([short_options, options,], 1)
            options = tf.strings.reduce_join(options, axis=-1, separator=": ")
            options = tf.strings.reduce_join(options, separator="\n")
            out["prompt"] = tf.strings.join([ex["question"], "\n", options, "\n"])
            if tf.reduce_sum(tf.cast(tf.strings.regex_full_match(options, "<img='(.*?)'>"), tf.int32)) > 1:
                # Following LLaVa, don't use any images if there are multiple images paths
                # I think the rationale is that this means the image are answer-options
                out["image"] = out["image"][:0]
        else:
            out["metadata/options"] = tf.constant([""] * 9, dtype=tf.string)
            out["prompt"] = ex["question"]
            out["image"] = out["image"][:0]
        return out
    ds = ds.map(_extract, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return ds

@gin.configurable()
def extract_mmmu_cot(ds, types=['multiple-choice', 'open'], styles=['ai2_diagram', 'vqa2'], default_style='ai2_diagram', option_format="abc"):
    assert option_format == "abc"
    keys_tensor = tf.constant(types, dtype=tf.string)
    values_tensor = tf.constant(styles, dtype=tf.string)
    table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(keys_tensor, values_tensor),
        default_value=tf.constant(default_style, dtype=tf.string),
    )
    def _extract(ex):
        # out = dict(image=tf.expand_dims(ex["image_with_question"], 0))
        out = dict(image=tf.expand_dims(ex["image_1"], 0))
        out.update(_add_metadata(ex))
        style = table.lookup(ex["metadata/question_type"])
        # out["style"] = style
        out["text"] = ex["answer"]
        out["metadata/question"] = ex["question"]
        out["metadata/references"] = ex["answer"]

        if style == styles[0]:
            abc = tf.constant(list("abcdefghi".upper()))
            options = ex["options"]
            num_options = tf.shape(options)[0]
            dummy_options = tf.tile(tf.constant([""], dtype=tf.string), [9 - num_options])
            out["metadata/options"] = tf.concat([options, dummy_options], axis=0)
            out["metadata/options"] = tf.ensure_shape(out["metadata/options"], [9])
            
            short_options = abc[:num_options]
            options = tf.stack([short_options, options,], 1)
            options = tf.strings.reduce_join(options, axis=-1, separator=": ")
            options = tf.strings.reduce_join(options, separator="\n")
            out["prompt"] = tf.strings.join([ex["question"], "\n", options, "\n"])
            # out["prompt"] = ex["question"]
            if tf.reduce_sum(tf.cast(tf.strings.regex_full_match(options, "<img='(.*?)'>"), tf.int32)) > 1:
                # Following LLaVa, don't use any images if there are multiple images paths
                # I think the rationale is that this means the image are answer-options
                out["image"] = out["image"][:0]
        else:
            out["metadata/options"] = tf.constant([""] * 9, dtype=tf.string)
            out["prompt"] = ex["question"]
            # out["image"] = out["image"][:0]
        return out
    ds = ds.map(_extract, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return ds


@seqio.map_over_dataset
def reformat_math_vista(ex):
    query = ex["query"]
    query = tf.strings.split(query, sep="Question:")[-1]
    query = tf.strings.strip(tf.strings.split(query, sep="Hint:")[0])
    ex["query"] = query
    return ex


@seqio.map_over_dataset
def extract_math_vista(ex, styles=['ai2_diagram', 'vqa2']):
    out = dict(image=ex["image"])
    out.update(_add_metadata(ex))

    is_mc = ex["metadata/question_type"] == 'multi_choice'
    if is_mc:
        style = styles[0]
        abc = tf.constant(list("abcdefghi".upper()))
        options = ex["choices"]
        num_options = tf.shape(options)[0]
        dummy_options = tf.tile(tf.constant([""], dtype=tf.string), [9 - num_options])
        out["metadata/options"] = tf.concat([options, dummy_options], axis=0)
        out["metadata/options"] = tf.ensure_shape(out["metadata/options"], [9])

        if ex["metadata/split"] != "test":
            short_options = abc[:num_options]
            answer_short_option = tf.boolean_mask(short_options, options == ex["answer"])[0]
            out["text"] = answer_short_option
        else:
            out["text"] = ex["answer"]
    else:
        style = styles[1]
        out["metadata/options"] = tf.constant([""] * 9, dtype=tf.string)
        out["text"] = ex["answer"]
    out["style"] = style
    out["prompt"] = ex["query"]
    out["metadata/query"] = ex["query"]
    out["metadata/references"] = ex["answer"]
    return out


NO_POINT_PREFIX = [
    "No pointing: ",
    "No pointing: ",
    "no pointing:\n",
    "No pointing:\n",
    "Not pointing:\n",
    "No Points: ",
    "No Points: ",
    "NO POINTING\n",
    "No pontiing\n",
    "No Points:\n ",
    "No pointing\n",
    "Do not point. ",
    "Refrain from pointing. ",
    "Avoid generating points . ",
    "For this question, do not use points. ",
    "Refrain from using points:\n",
    "Don't include points in your response. ",
    "Don't point. ",
    "Don't use points. ",
    "Please don't use points.\n\n",
    "Please don't use points.\n\n",
    "Respond without using points. ",
    "Respond without pointing:\n",
    "Do not generate ponits: ",
    "Do not point. ",
    "Do not point\n",
    "no pointing\n\n",
    "Answer without points: ",
    "Answer this question without pointing: ",
    "Answer without poiints. ",
    "answer without points: ",
    "answer with text only, do not points\n"
]
assert all(x[-1].isspace() for x in NO_POINT_PREFIX)
NO_POINT_PREFIX_TF = tf.constant(NO_POINT_PREFIX)


def prefix_how_many(messages, seed):
    question = messages[0]
    if tf.strings.regex_full_match(tf.strings.lower(question), "how many.*"):
        ix = tf.random.stateless_uniform((), seed, 0, len(NO_POINT_PREFIX),  tf.int32)
        question = tf.strings.join([NO_POINT_PREFIX_TF[ix], question])
        return tf.concat([tf.expand_dims(question, 0), messages[1:]], axis=0)
    else:
        return messages


@seqio.map_over_dataset(num_seeds=1)
def prefix_how_many_messages(ex, seed):
    messages = ex["messages"]
    n = tf.shape(messages)[0]
    seeds = tf.random.split(seed, n)
    message_arr = tf.TensorArray(dtype=tf.string, size=n, element_shape=(None,))
    for i in range(n):
        message_arr = message_arr.write(i, prefix_how_many(messages[i], seeds[i]))
    ex["messages"] = tf.RaggedTensor.from_row_splits(
        values=message_arr.concat(), row_splits=messages.row_splits)
    return ex


def filter_single_turn(ds):
    @seqio.map_over_dataset
    def _filter(ex):
        multi_turn = ex["messages"].row_lengths() > 2
        ex["messages"] = tf.ragged.boolean_mask(ex["messages"], multi_turn)
        return ex

    ds = _filter(ds)
    ds = ds.filter(lambda x: tf.shape(x["messages"])[0] > 0)
    return ds


@seqio.map_over_dataset(num_seeds=1)
def extract_cockatoo_qa_v2(ex, seed):
    messages = tf.RaggedTensor.from_value_rowids(ex["messages"], ex["conversation_ids"])
    ix = stateless_permutation(tf.shape(messages)[0], seed)
    messages = tf.gather(messages, ix)
    out = dict(
        image=ex["image"],
        messages=messages
    )
    out.update(_add_metadata(ex))
    return out


def format_mmbench(ds):

    def _trim(ex):
        num_passes = tf.shape(ex["id"])[0]
        ex["choices"] = ex["choices"][:num_passes, :num_passes]
        ex["answer"] = ex["answer"][:num_passes]
        return ex

    ds = ds.map(_trim)
    ds = flatten_parts(ds, ["id", "query", "choices", "answer"])

    def _extract(ex):
        out = dict(image=ex["image"])
        out.update(_add_metadata(ex))
        out["prompt"] = ex["query"]
        out["text"] = ex["answer"]
        options = ex["choices"]
        tf.debugging.assert_equal(tf.reduce_any(tf.strings.regex_full_match(options, ".*\|\|\|.*")), False)
        out["metadata/options"] = tf.strings.reduce_join(options, separator="|||")
        out["metadata/question"] = ex["question"]
        out["metadata/references"] = ex["answer"]
        return out

    ds = ds.map(_extract, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return ds


@seqio.map_over_dataset
def extract_lvis(ex, class_name_file="gs://oe-training-chrisc/cockatoo/data/lvis_class_names.json"):
    with tf.io.gfile.GFile(class_name_file) as f:
        class_names = json.load(f)
    class_names_arr = [None]*len(class_names)
    for k, v in class_names.items():
        class_names_arr[int(k)] = v
    assert all(x is not None for x in class_names_arr)
    class_names_arr = tf.constant(class_names_arr)

    return dict(
        image=ex["image"],
        bbox=ex["objects"]["bbox"],
        label=tf.gather(class_names_arr, ex["objects"]["label"]),
    )


def extract_open_images_boxes(ds):
    # ds = ds.filter(lambda ex: tf.logical_or(
    #     tf.shape(ex["cap/cap_caption"])[0] > 0,
    # tf.shape(ex["detection/bbox"])[0] > 0
    # ))
    ds = ds.filter(lambda ex: tf.shape(ex["cap/cap_caption"])[0] > 0)

    @seqio.map_over_dataset
    def _map(ex):
        bbox = tf.reshape(ex["detection/bbox"], (-1, 4))
        bbox = tf.stack([
            bbox[:, 2],
            bbox[:, 0],
            bbox[:, 3],
            bbox[:, 1]
        ], 1)
        return dict(
            image=tf.image.decode_jpeg(ex["image"]),
            bbox=bbox,
            label=ex["detection/label"],
            caption=tf.strings.reduce_join(ex["cap/cap_caption"], separator="\n")
        )

    return _map(ds)


@seqio.map_over_dataset
def region_captions_to_dense(ex):
    if "captions" in ex:
        captions = ex["captions"]["text"]
        boxes = ex["captions"]["bbox"]
    else:
        captions = ex["label"]
        boxes = ex["bbox"]


    sh = tf.cast(tf.shape(ex["image"])[:2], tf.float32)
    # image_h, image_w = sh[0], sh[1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]

    cx = tf.cast(boxes[:, 0] + w/2, tf.float32)
    cy = tf.cast(boxes[:, 1] + h/2, tf.float32)
    # w = w / image_w
    # h = h / image_h
    coor = tf.strings.reduce_join(
        float_to_text(tf.stack([cx, cy, w, h], 1)), separator=",", axis=1)

    area = w*h
    if tf.random.uniform(()) < 0.5:
        coor_text = "before"
        captions = tf.strings.join([coor, captions], separator=": ")
    else:
        coor_text = "after"
        captions = tf.strings.join([captions, coor], separator=": ")

    ix = tf.random.uniform((), 0, 6, tf.int32)
    center = boxes
    if ix == 0:
        order_text = "left"
        sort_by = boxes[:, 0]
    elif ix == 1:
        order_text = "right"
        sort_by = -boxes[:, 2]
    elif ix == 2:
        order_text = "top"
        sort_by = boxes[:, 1]
    elif ix == 3:
        order_text = "bottom"
        sort_by = -boxes[:, 3]
    elif ix == 4:
        order_text = "largest"
        sort_by = area
    else:
        order_text = "smallest"
        sort_by = -area
    ixs = tf.argsort(sort_by)
    captions = tf.gather(captions, ixs)
    text = tf.strings.join([
        order_text,
        coor_text,
        tf.strings.reduce_join(captions, separator="\n")
    ], separator="; ")

    if "caption" in ex:
        if tf.random.uniform(()) > 0.5:
            text = tf.strings.join([text, "\ncaption: ", ex["caption"]])
        else:
            text = tf.strings.join(["caption: ", ex["caption"], "\n", text])

    return dict(
        image=ex["image"],
        text=text
    )


@seqio.map_over_dataset()
def join_captions(ex):
    text = tf.random.shuffle(ex['text'])
    ex["text"] = tf.strings.reduce_join(text, separator="\n")
    return ex


@seqio.map_over_dataset(num_seeds=1)
def extract_figureqa(ex, seed):
    questions = ex["questions"]
    n = stateless_permutation(tf.shape(questions["question"])[0], seed)
    return dict(
        image=ex["image"],
        questions=tf.gather(questions["question"], n),
        question_id=tf.gather(questions["question_id"], n),
        answer=tf.gather(tf.strings.as_string(questions["answer"]), n)
    )


@seqio.map_over_dataset
def convert_figureqa_answer(ex):
    keys_tensor = tf.constant(["0", "1"])
    values_tensor = tf.constant(["no", "yes"])
    table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(keys_tensor, values_tensor),
        default_value=tf.constant("nan", dtype=tf.string),
    )
    answer = table.lookup(ex["answer"])
    ex["answer"] = answer
    return ex


@seqio.map_over_dataset()
def build_question_with_hint(ex):
    hint = ex["hint"]
    if tf.strings.length(hint) > 0:
        ex["question"] = tf.strings.join([hint, ex["question"]], separator="\n")
    return ex

@seqio.map_over_dataset()
def build_question_with_context(ex):
    context = ex["context"]
    if tf.strings.length(context) > 0:
        ex["question"] = tf.strings.join([context, ex["question"]], separator="\n")
    return ex


def max_words(ds, max_words):
    return ds.filter(lambda x: x["n_words"] <= max_words)


@seqio.map_over_dataset
def format_pdfa_eng_wds(example):
    return dict(
        image=example["image"],
        text=tf.strings.reduce_join(example["lines"]["text"], separator="\n"),
    )


@gin.configurable()
def accuracy_conditioned_joint(ds, sequence_length, is_eval=False, eval_quality=17,
                               transcript_quality=None):
    # v2: Transcripts no longer get a quality score
    is_training = sequence_length.get('is_training', True)
    if not is_training:
        if is_eval:
            prompt = f"quality {eval_quality}:"
        else:
            prompt = f"quality 17:"

        @seqio.map_over_dataset
        def _with_prompt(ex):
            out = dict(
                image=ex["image"],
                url=ex["url"],
                prompt=prompt,
            )
            if "text" in ex:
                out["text"] = ex["text"]
            elif "caption" in ex:
                out["text"] = ex["caption"]
            return out
        return _with_prompt(ds)

    elif is_eval:
        raise ValueError("is_eval=True and is_training=False")

    # each transcript
    @seqio.map_over_dataset
    def _with_transcript(ex):
        if tf.shape(ex["edited_captions"]["caption"])[0] > 0:
            edited_caption = ex["edited_captions"]["caption"][0]
            n = ex["edited_captions"]["n_edits"][0]
        else:
            edited_caption = ""
            n = 0
        text = [
            ex["caption"],
            ex["transcripts"][tf.random.uniform((), 0, tf.shape(ex["transcripts"])[0], dtype=tf.int32)],
            edited_caption
        ]
        edit_quality = 17 - n
        prompt = [
            "quality 17:",
            "" if transcript_quality is None else f"quality: {edit_quality}:",
            tf.strings.join(["quality ", tf.strings.as_string(edit_quality), ":"])
        ]
        return dict(
            image=ex["image"],
            text=tf.stack(text, 0),
            url=ex["url"],
            prompt=tf.stack(prompt, 0),
            style=["long_caption", "transcript", "long_caption"]
        )
    return _with_transcript(ds)


def select_dense_caption_sample(ds, samples=200):
    def compute_hash(string: str) -> str:
        return hashlib.sha256(string.encode("utf-8")).hexdigest()

    with tf.io.gfile.GFile("gs://oe-training-chrisc/cockatoo/data/dense-caption-eval-v0-final-data.json") as f:
        data = json.load(f)
    for ex in data:
        ex["image_id"] = compute_hash(ex["image"])
    data.sort(key=lambda x: x["image_id"])
    np.random.RandomState(12312).shuffle(data)
    keep = tf.constant([x["image"] for x in data[:samples]])

    def _keep(ex):
        return tf.reduce_any(ex["url"] == keep)
    ds = ds.filter(_keep)
    ds = tf.data.experimental.assert_cardinality(samples)(ds)
    return ds

@seqio.map_over_dataset()
def charxiv_preprocessor(ex):
    question_names = ["descriptive_q1", "descriptive_q2", "descriptive_q3", "descriptive_q4", "reasoning_q"]
    answer_names = ["descriptive_a1", "descriptive_a2", "descriptive_a3", "descriptive_a4", "reasoning_a"]

    questions = [ex[name] for name in question_names]
    answers = [ex[name] for name in answer_names]

    return dict(
        image=ex["image"],
        question=tf.stack(questions, 0),
        answer=tf.stack(answers, 0)
    )

@seqio.map_over_dataset()
def charxiv_descriptive_preprocessor(ex):
    question_names = ["descriptive_q1", "descriptive_q2", "descriptive_q3", "descriptive_q4"]
    answer_names = ["descriptive_a1", "descriptive_a2", "descriptive_a3", "descriptive_a4"]

    questions = [ex[name] for name in question_names]
    answers = [ex[name] for name in answer_names]

    return dict(
        image=ex["image"],
        question=tf.stack(questions, 0),
        answer=tf.stack(answers, 0)
    )

@seqio.map_over_dataset()
def charxiv_reasoning_preprocessor(ex):
    return dict(
        image=ex["image"],
        question=ex["reasoning_q"],
        answer=ex["reasoning_a"]
    )

@seqio.map_over_dataset()
def tablevqa_preprocessor(ex):
    return dict(
        image=ex["image"],
        question=ex["question"],
        answer=ex["gt"]
    )

@seqio.map_over_dataset()
def vtabfact_preprocessor(ex):
    return dict(
        image=ex["image"],
        question=tf.strings.join([ex["question"], "Answer with yes or no."], separator="\n"),
        answer=ex["gt"]
    )

@seqio.map_over_dataset()
def nutrition_fact_preprocessor(ex):
    question_names = ["descriptive_q", "reasoning_q"]
    answer_names = ["descriptive_a", "reasoning_a"]

    questions = [ex[name] for name in question_names]
    answers = [ex[name] for name in answer_names]

    return dict(
        image=ex["image"],
        question=tf.stack(questions, 0),
        answer=tf.stack(answers, 0)
    )