src/StructDiffusion/language/tokenizer.py

import json
import numpy as np
import re

# def add_pad_to_vocab(vocab):
#     new_vocab = {"PAD": 0}
#     for k in vocab:
#         new_vocab[k] = vocab[k] + 1
#     return new_vocab
#
#
# def combine_vocabs(vocabs, vocab_types):
#     new_vocab = {}
#     for type, vocab in zip(vocab_types, vocabs):
#         for k in vocab:
#             new_vocab["{}:{}".format(type, k)] = len(new_vocab)
#     return new_vocab
#
#
# def add_token_to_vocab(vocab):
#     new_vocab = {"MASK": 0}
#     for k in vocab:
#         new_vocab[k] = vocab[k] + 1
#     return new_vocab
#
#
# def tokenize_circle_specification(circle_specification):
#     tokenized = {}
#     # min 0, max 0.5, increment 0.05, 10 discrete values
#     tokenized["radius"] = int(circle_specification["radius"] / 0.05)
#
#     # min 0, max 1, increment 0.10, 10 discrete values
#     tokenized["position_x"] = int(circle_specification["position"][0] / 0.10)
#
#     # min -0.5, max 0.5, increment 0.10, 10 discrete values
#     tokenized["position_y"] = int(circle_specification["position"][1] / 0.10)
#
#     # min -3.14, max 3.14, increment 3.14 / 18, 36 discrete values
#     tokenized["rotation"] = int((circle_specification["rotation"][2] + 3.14) / (3.14 / 18))
#
#     uniform_angle_vocab = {"False": 0, "True": 1}
#     tokenized["uniform_angle"] = uniform_angle_vocab[circle_specification["uniform_angle"]]
#
#     face_center_vocab = {"False": 0, "True": 1}
#     tokenized["face_center"] = face_center_vocab[circle_specification["face_center"]]
#
#     angle_ratio_vocab = {0.5: 0, 1.0: 1}
#     tokenized["angle_ratio"] = angle_ratio_vocab[circle_specification["angle_ratio"]]
#
#     # heights min 0.0, max 0.5
#     # volumn min 0.0, max 0.012
#
#     return tokenized
#
#
# def build_vocab(old_vocab_file, new_vocab_file):
#     with open(old_vocab_file, "r") as fh:
#         vocab_json = json.load(fh)
#
#     vocabs = {}
#     vocabs["class"] = vocab_json["class_to_idx"]
#     vocabs["size"] = vocab_json["size_to_idx"]
#     vocabs["color"] = vocab_json["color_to_idx"]
#     vocabs["material"] = vocab_json["material_to_idx"]
#     vocabs["comparator"] = {"less": 1, "greater": 2, "equal": 3}
#
#     vocabs["radius"] = (0.0, 0.5, 10)
#     vocabs["position_x"] = (0.0, 1.0, 10)
#     vocabs["position_y"] = (-0.5, 0.5, 10)
#     vocabs["rotation"] = (-3.14, 3.14, 36)
#     vocabs["height"] = (0.0, 0.5, 10)
#     vocabs["volumn"] = (0.0, 0.012, 10)
#
#     vocabs["uniform_angle"] = {"False": 0, "True": 1}
#     vocabs["face_center"] = {"False": 0, "True": 1}
#     vocabs["angle_ratio"] = {0.5: 0, 1.0: 1}
#
#     with open(new_vocab_file, "w") as fh:
#         json.dump(vocabs, fh)


class Tokenizer:
    """
    We want to build a tokenizer that tokenize words, features, and numbers.

    This tokenizer should also allow us to sample random values.

    For discrete values, we store mapping from the value to an id
    For continuous values, we store min, max, and number of bins after discretization

    """

    def __init__(self, vocab_file):

        self.vocab_file = vocab_file
        with open(self.vocab_file, "r") as fh:
            self.type_vocabs = json.load(fh)

        self.vocab = {"PAD": 0, "CLS": 1}
        self.discrete_types = set()
        self.continuous_types = set()
        self.build_one_vocab()

        self.object_position_vocabs = {}
        self.build_object_position_vocabs()

    def build_one_vocab(self):
        print("\nBuild one vacab for everything...")

        for typ, vocab in self.type_vocabs.items():
            if typ == "comparator":
                continue

            if typ in ["obj_x", "obj_y", "obj_z", "obj_rr", "obj_rp", "obj_ry",
                       "struct_x", "struct_y", "struct_z", "struct_rr", "struct_rp", "struct_ry"]:
                continue

            if type(vocab) == dict:
                self.vocab["{}:{}".format(typ, "MASK")] = len(self.vocab)

                for v in vocab:
                    assert ":" not in v
                    self.vocab["{}:{}".format(typ, v)] = len(self.vocab)
                self.discrete_types.add(typ)

            elif type(vocab) == tuple or type(vocab) == list:
                self.vocab["{}:{}".format(typ, "MASK")] = len(self.vocab)

                for c in self.type_vocabs["comparator"]:
                    self.vocab["{}:{}".format(typ, c)] = len(self.vocab)

                min_value, max_value, num_bins = vocab
                for i in range(num_bins):
                    self.vocab["{}:{}".format(typ, i)] = len(self.vocab)
                self.continuous_types.add(typ)
            else:
                raise TypeError("The dtype of the vocab cannot be handled: {}".format(vocab))

        print("The vocab has {} tokens: {}".format(len(self.vocab), self.vocab))

    def build_object_position_vocabs(self):
        print("\nBuild vocabs for object position")
        for typ in ["obj_x", "obj_y", "obj_z", "obj_rr", "obj_rp", "obj_ry",
                    "struct_x", "struct_y", "struct_z", "struct_rr", "struct_rp", "struct_ry"]:
            self.object_position_vocabs[typ] = {"PAD": 0, "MASK": 1}

            if typ not in self.type_vocabs:
                continue
            min_value, max_value, num_bins = self.type_vocabs[typ]
            for i in range(num_bins):
                self.object_position_vocabs[typ]["{}".format(i)] = len(self.object_position_vocabs[typ])
            print("The {} vocab has {} tokens: {}".format(typ, len(self.object_position_vocabs[typ]), self.object_position_vocabs[typ]))

    def get_object_position_vocab_sizes(self):
        return len(self.object_position_vocabs["position_x"]), len(self.object_position_vocabs["position_y"]), len(self.object_position_vocabs["rotation"])

    def get_vocab_size(self):
        return len(self.vocab)

    def tokenize_object_position(self, value, typ):
        assert typ in ["obj_x", "obj_y", "obj_z", "obj_rr", "obj_rp", "obj_ry",
                       "struct_x", "struct_y", "struct_z", "struct_rr", "struct_rp", "struct_ry"]
        if value == "MASK" or value == "PAD":
            return self.object_position_vocabs[typ][value]
        elif value == "IGNORE":
            # Important: used to avoid computing loss. -100 is the default ignore_index for NLLLoss
            return -100
        else:
            min_value, max_value, num_bins = self.type_vocabs[typ]
            assert min_value <= value <= max_value, value
            dv = min(int((value - min_value) / ((max_value - min_value) / num_bins)), num_bins - 1)
            return self.object_position_vocabs[typ]["{}".format(dv)]

    def tokenize(self, value, typ=None):
        if value in ["PAD", "CLS"]:
            idx = self.vocab[value]
        else:
            if typ is None:
                raise KeyError("Type cannot be None")

            if typ[-2:] == "_c" or typ[-2:] == "_d":
                typ = typ[:-2]

            if typ in self.discrete_types:
                idx = self.vocab["{}:{}".format(typ, value)]
            elif typ in self.continuous_types:
                if value == "MASK" or value in self.type_vocabs["comparator"]:
                    idx = self.vocab["{}:{}".format(typ, "MASK")]
                else:
                    min_value, max_value, num_bins = self.type_vocabs[typ]
                    assert min_value <= value <= max_value, "type {} value {} exceeds {} and {}".format(typ, value, min_value, max_value)
                    dv = min(int((value - min_value) / ((max_value - min_value) / num_bins)), num_bins - 1)
                    # print(value, dv, "{}:{}".format(typ, dv))
                    idx = self.vocab["{}:{}".format(typ, dv)]
            else:
                raise KeyError("Do not recognize the type {} of the given token: {}".format(typ, value))
        return idx

    def get_valid_random_value(self, typ):
        """
        Get a random value for the given typ
        :param typ:
        :return:
        """
        if typ[-2:] == "_c" or typ[-2:] == "_d":
            typ = typ[-2:]

        candidate_values = []
        for v in self.vocab:
            if v in ["PAD", "CLS"]:
                continue
            ft, fv = v.split(":")
            if typ == ft and fv != "MASK" and fv not in self.type_vocabs["comparator"]:
                candidate_values.append(v)
        assert len(candidate_values) != 0
        typed_v = np.random.choice(candidate_values)
        value = typed_v.split(":")[1]

        if typ in self.discrete_types:
            return value
        elif typ in self.continuous_types:
            min_value, max_value, num_bins = self.type_vocabs[typ]
            return min_value + ((max_value - min_value) / num_bins) * int(value)
        else:
            raise KeyError("Do not recognize the type {} of the given token".format(typ))

    def get_all_values_of_type(self, typ):
        """
        Get all values for the given typ
        :param typ:
        :return:
        """
        if typ[-2:] == "_c" or typ[-2:] == "_d":
            typ = typ[-2:]

        candidate_values = []
        for v in self.vocab:
            if v in ["PAD", "CLS"]:
                continue
            ft, fv = v.split(":")
            if typ == ft and fv != "MASK" and fv not in self.type_vocabs["comparator"]:
                candidate_values.append(v)
        assert len(candidate_values) != 0
        values = [typed_v.split(":")[1] for typed_v in candidate_values]

        if typ in self.discrete_types:
            return values
        else:
            raise KeyError("Do not recognize the type {} of the given token".format(typ))

    def convert_to_natural_sentence(self, template_sentence):

        # select objects that are [red, metal]
        # select objects that are [larger, taller] than the [], [], [] object
        # select objects that have the same [color, material] of the [], [], [] object

        natural_sentence_templates = ["select objects that are {}.",
                                      "select objects that have {} {} {} the {}.",
                                      "select objects that have the same {} as the {}."]

        v, t = template_sentence[0]
        if t[-2:] == "_c" or t[-2:] == "_d":
            t = t[:-2]

        if v != "MASK" and t in self.discrete_types:
            natural_sentence_template = natural_sentence_templates[0]
            if t == "class":
                natural_sentence = natural_sentence_template.format(re.findall(r'[A-Z](?:[a-z]+|[A-Z]*(?=[A-Z]|$))', v)[0].lower())
            else:
                natural_sentence = natural_sentence_template.format(v)
        else:
            anchor_obj_properties = []
            class_reference = None
            for token in template_sentence[1:]:
                if token[0] != "PAD":
                    if token[1] == "class":
                        class_reference = token[0]
                    else:
                        anchor_obj_properties.append(token[0])
            # order the properties
            anchor_obj_des = ", ".join(anchor_obj_properties)
            if class_reference is None:
                anchor_obj_des += " object"
            else:
                anchor_obj_des += " {}".format(re.findall(r'[A-Z](?:[a-z]+|[A-Z]*(?=[A-Z]|$))', class_reference)[0].lower())

            if v == "MASK":
                natural_sentence_template = natural_sentence_templates[2]
                anchor_type = t
                natural_sentence = natural_sentence_template.format(anchor_type, anchor_obj_des)
            elif t in self.continuous_types:
                natural_sentence_template = natural_sentence_templates[1]
                if v == "equal":
                    jun = "as"
                else:
                    jun = "than"
                natural_sentence = natural_sentence_template.format(v, t, jun, anchor_obj_des)
            else:
                raise NotImplementedError

        return natural_sentence

    def prepare_grounding_reference(self):
        goal = {"rearrange": {"features": []},
                "anchor": {"features": []}}
        discrete_type = ["class", "material", "color"]
        continuous_type = ["volumn", "height"]

        print("#"*50)
        print("Preparing referring expression")

        refer_type = verify_input("direct (1) or relational reference (2)? ", [1, 2], int)
        if refer_type == 1:

            # 1. no anchor
            t = verify_input("desired type: ", discrete_type, None)
            v = verify_input("desired value: ", self.get_all_values_of_type(t), None)

            goal["rearrange"]["features"].append({"comparator": None, "type": t, "value": v})

        elif refer_type == 2:

            value_type = verify_input("discrete (1) or continuous relational reference (2)? ", [1, 2], int)
            if value_type == 1:
                t = verify_input("desired type: ", discrete_type, None)
                # 2. discrete
                goal["rearrange"]["features"].append({"comparator": None, "type": t, "value": None})
            elif value_type == 2:
                comp = verify_input("desired comparator: ", list(self.type_vocabs["comparator"].keys()), None)
                t = verify_input("desired type: ", continuous_type, None)
                # 3. continuous
                goal["rearrange"]["features"].append({"comparator": comp, "type": t, "value": None})

            num_f = verify_input("desired number of features for the anchor object: ", [1, 2, 3], int)
            for i in range(num_f):
                t = verify_input("desired type: ", discrete_type, None)
                v = verify_input("desired value: ", self.get_all_values_of_type(t), None)
                goal["anchor"]["features"].append({"comparator": None, "type": t, "value": v})

        return goal

    def convert_structure_params_to_natural_language(self, sentence):

        # ('circle', 'shape'), (-1.3430555575431449, 'rotation'), (0.3272675147405848, 'position_x'), (-0.03104362197706456, 'position_y'), (0.04674859577847633, 'radius')

        shape = None
        x = None
        y = None
        rot = None
        size = None

        for param in sentence:
            if param[0] == "PAD":
                continue

            v, t = param
            if t == "shape":
                shape = v
            elif t == "position_x":
                dv = self.discretize(v, t)
                if dv == 0:
                    x = "bottom"
                elif dv == 1:
                    x = "middle"
                elif dv == 2:
                    x = "top"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))
            elif t == "position_y":
                dv = self.discretize(v, t)
                if dv == 0:
                    y = "right"
                elif dv == 1:
                    y = "center"
                elif dv == 2:
                    y = "left"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))
            elif t == "radius":
                dv = self.discretize(v, t)
                if dv == 0:
                    size = "small"
                elif dv == 1:
                    size = "medium"
                elif dv == 2:
                    size = "large"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))
            elif t == "rotation":
                dv = self.discretize(v, t)
                if dv == 0:
                    rot = "north"
                elif dv == 1:
                    rot = "east"
                elif dv == 2:
                    rot = "south"
                elif dv == 3:
                    rot = "west"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))

        natural_sentence = "" # "{} {} in the {} {} of the table facing {}".format(size, shape, x, y, rot)

        if size:
            natural_sentence += "{}".format(size)
        if shape:
            natural_sentence += " {}".format(shape)
        if x:
            natural_sentence += " in the {}".format(x)
        if y:
            natural_sentence += " {} of the table".format(y)
        if rot:
            natural_sentence += " facing {}".format(rot)

        natural_sentence = natural_sentence.strip()

        return natural_sentence

    def convert_structure_params_to_type_value_tuple(self, sentence):

        # ('circle', 'shape'), (-1.3430555575431449, 'rotation'), (0.3272675147405848, 'position_x'), (-0.03104362197706456, 'position_y'), (0.04674859577847633, 'radius')

        shape = None
        x = None
        y = None
        rot = None
        size = None

        for param in sentence:
            if param[0] == "PAD":
                continue

            v, t = param
            if t == "shape":
                shape = v
            elif t == "position_x":
                dv = self.discretize(v, t)
                if dv == 0:
                    x = "bottom"
                elif dv == 1:
                    x = "middle"
                elif dv == 2:
                    x = "top"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))
            elif t == "position_y":
                dv = self.discretize(v, t)
                if dv == 0:
                    y = "right"
                elif dv == 1:
                    y = "center"
                elif dv == 2:
                    y = "left"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))
            elif t == "radius":
                dv = self.discretize(v, t)
                if dv == 0:
                    size = "small"
                elif dv == 1:
                    size = "medium"
                elif dv == 2:
                    size = "large"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))
            elif t == "rotation":
                dv = self.discretize(v, t)
                if dv == 0:
                    rot = "north"
                elif dv == 1:
                    rot = "east"
                elif dv == 2:
                    rot = "south"
                elif dv == 3:
                    rot = "west"
                else:
                    raise KeyError("key {} not found in {}".format(v, self.type_vocabs[t]))

        # rotation, shape, size, x, y
        type_value_tuple_init = [("rotation", rot), ("shape", shape), ("size", size), ("x", x), ("y", y)]
        type_value_tuple = []
        for type_value in type_value_tuple_init:
            if type_value[1] is not None:
                type_value_tuple.append(type_value)

        type_value_tuple = tuple(sorted(type_value_tuple))
        return type_value_tuple

    def discretize(self, v, t):
        min_value, max_value, num_bins = self.type_vocabs[t]
        assert min_value <= v <= max_value, "type {} value {} exceeds {} and {}".format(t, v, min_value, max_value)
        dv = min(int((v - min_value) / ((max_value - min_value) / num_bins)), num_bins - 1)
        return dv


class ContinuousTokenizer:
    """
    This tokenizer is for testing not discretizing structure parameters
    """

    def __init__(self):

        print("WARNING: Current continous tokenizer does not support multiple shapes")

        self.continuous_types = ["rotation", "position_x", "position_y", "radius"]
        self.discrete_types = ["shape"]

    def tokenize(self, value, typ=None):
        if value == "PAD":
            idx = 0.0
        else:
            if typ is None:
                raise KeyError("Type cannot be None")
            elif typ in self.discrete_types:
                idx = 1.0
            elif typ in self.continuous_types:
                idx = value
            else:
                raise KeyError("Do not recognize the type {} of the given token: {}".format(typ, value))
        return idx