lt_s2s_sup_tcd.py

"""
Infer trained model with type constrained decoding.
"""
import os
import argparse
import logging
from pathlib import Path
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# from dataset_lifted import load_split_dataset
# from eval import evaluate_sym_trans
# from utils import count_params


PROPS = ["a", "b", "c", "d", "h"]
T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b"]
T5_PREFIX = "translate English to Linear Temporal Logic: "


S2S_MODELS = T5_MODELS.extend(["pt_transformer"])
UNARY_OPERATORS = ['!', "F", "G", "X"]
BINARY_OPERATORS = ['&', '|', 'U', 'i', 'e', 'M']
END_TOKEN = '</s>'
MAX_LENGTH = 256  # 266? 2 x exact_restricted_avoidance_5
MAX_DEPTH = 21
CHECK_DEPTH = 100


def is_valid(formula: list[str], next_token):
    if len(formula) == 0:
        return True
    else:
        if formula[0] in UNARY_OPERATORS:
            prop_counter = 1
        elif formula[0] in BINARY_OPERATORS:
            prop_counter = 2
        else:
            prop_counter = 0

        for i in range(1, len(formula)):
            if formula[i] not in UNARY_OPERATORS and formula[i] not in BINARY_OPERATORS:
                prop_counter -= 1
            elif formula[i] in BINARY_OPERATORS:
                prop_counter += 1

        if next_token in UNARY_OPERATORS:
            return prop_counter > 0
        elif next_token in BINARY_OPERATORS:
            return prop_counter > 0
        elif next_token == END_TOKEN:
            return prop_counter == 0
        else:
            return prop_counter > 0


class Node:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


def depth(root):
    """
    Returns the depth of a binary tree
    """
    if root is None:
        return 0
    leftAns = depth(root.left)
    rightAns = depth(root.right)
    return max(leftAns, rightAns) + 1


def build_tree(ltl):
    """
    Construct a binary tree from LTL formula string and return as p, partial formula supported
    """
    if ltl == '':  # its the end of the expression
        return Node(''), ''
    if ltl[0] in PROPS:  # the character is a proposition
        return Node(ltl[0]), ltl[1:]
    elif ltl[0] in BINARY_OPERATORS:
        # Create a node with ltl[0] as the data and both the children set to null
        p = Node(ltl[0])
        # Build the left sub-tree
        p.left, q = build_tree(ltl[1:])
        # Build the right sub-tree
        p.right, q = build_tree(q)
        return p, q


class Seq2Seq:
    def __init__(self, model_dpath, model_name, **kwargs):
        self.model_name = model_name
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if "t5" in model_name or "bart" in model_name:
            self.tokenizer = AutoTokenizer.from_pretrained(model_dpath)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dpath).to(self.device)
        else:
            raise ValueError(f'ERROR: unrecognized model: {model_name}')

    def translate(self, queries):
        if "t5" in self.model_name or "bart" in self.model_name:
            inputs = [f"{T5_PREFIX}{query}" for query in queries]  # add prefix

            inputs = self.tokenizer(inputs, return_tensors="pt", padding=True).to(self.device)
            output_tokens = self.model.generate(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
                do_sample=False,
                max_new_tokens=256,
            )
            ltls = self.tokenizer.batch_decode(output_tokens, skip_special_tokens=True)
        elif self.model_name == "pt_transformer":
            ltls = [self.model_translate(self.model, self.vocab_transform, self.text_transform, queries[0])]
        else:
            raise ValueError(f'ERROR: unrecognized model, {self.model_name}')
        return ltls

    def type_constrained_decode(self, utts):
        """
        type constrained decoding based on LTL syntax:
        mask_and_regen() :: next_token is invalid, mask it and pick the token with the second highest likelihood
        add_and_gen_new() :: next_token is valid, append it to the partial formula and generate a new one

        Logic:
        if len(ltl) == 0: mask_and_regen(next_token) if next_token == whitespace
            |
        elif len(ltl) > CHECK_DEPTH: if depth(ltl) > MAX_DEPTH: only props and <EOS> allowed
            |                                       |
            |                                       else: regular_generation(next_token)
            |
        elif not next_token == whitespace: regular_generation(next_token)
            |
        elif next_token == whitespace: mask_and_regen(next_token) if ltl[-1] == whitespace else add_and_gen_new(next_token)
            |
        else: mask_and_regen(next_token)
        """
        def mask_and_regen(lm_logits, next_decoder_input_ids):
            lm_logits[:, -1:, next_decoder_input_ids.item()] = float('-inf')
            next_decoder_input_ids = torch.argmax(lm_logits[:, -1:], axis=-1)
            next_token = self.tokenizer.decode(next_decoder_input_ids[0], skip_special_tokens=False)
            return next_token, next_decoder_input_ids, lm_logits

        def add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace):
            whitespace = next_token == ''
            decoder_input_ids = torch.cat([decoder_input_ids, next_decoder_input_ids], axis=-1)
            lm_logits = self.model(None, encoder_outputs=encoded_sequence, decoder_input_ids=decoder_input_ids, return_dict=True).logits
            next_decoder_input_ids = torch.argmax(lm_logits[:, -1:], axis=-1)
            next_token = self.tokenizer.decode(next_decoder_input_ids[0], skip_special_tokens=False)
            return next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace

        if "t5" in self.model_name or "bart" in self.model_name:
            inputs = [f"{T5_PREFIX}{utt}" for utt in utts]  # add prefix
            inputs = self.tokenizer(inputs, return_tensors="pt", padding=True).to(self.device)
            input_ids = inputs.input_ids
            decoder_input_ids = self.tokenizer("<pad>", add_special_tokens=False, return_tensors="pt").input_ids.to(self.device)
            outputs = self.model(input_ids, decoder_input_ids=decoder_input_ids, return_dict=True)
            encoded_sequence = (outputs.encoder_last_hidden_state,)
            lm_logits = outputs.logits
            next_decoder_input_ids = torch.argmax(lm_logits[:, -1:], axis=-1)
            next_token = self.tokenizer.decode(next_decoder_input_ids[0], skip_special_tokens=False)
            token_list = []
            whitespace = False
            while '</s>' not in token_list and len(token_list) < MAX_LENGTH:
                # decode the first token
                if len(token_list) == 0:
                    if next_token in PROPS + UNARY_OPERATORS + BINARY_OPERATORS:
                        token_list.append(next_token)
                        next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                    elif next_token == '':
                        next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                    else:
                        next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)
                # after the first token
                # only check depth after certain # of operators & props
                elif len(token_list) > CHECK_DEPTH:
                    no_uni_list = ''.join([s for s in token_list if not s in UNARY_OPERATORS])
                    partial_tree, _ = build_tree(no_uni_list)
                    # if max_depth is reached
                    if depth(partial_tree) > MAX_DEPTH:
                        # only output props and <EOS>
                        if not (next_token in PROPS or next_token == END_TOKEN or next_token == ''):
                            next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)
                        elif next_token in PROPS or next_token == END_TOKEN:
                            if not is_valid(token_list, next_token):
                                next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)

                            else:
                                token_list.append(next_token)
                                next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                        # no consecutive whitespaces
                        elif next_token == ' ' or next_token == '':
                            if whitespace == True:
                                next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)
                            else:
                                next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                        # mask all other tokens
                        else:
                            next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)
                    else:
                        if next_token in PROPS + UNARY_OPERATORS + BINARY_OPERATORS + [END_TOKEN]:
                            if not is_valid(token_list, next_token):
                                next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)

                            else:
                                token_list.append(next_token)
                                next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                        # no consecutive whitespaces
                        elif next_token == ' ' or next_token == '':
                            if whitespace == True:
                                next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)
                            else:
                                next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                # mask all other tokens
                        else:
                            next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)

                elif next_token in PROPS + UNARY_OPERATORS + BINARY_OPERATORS + [END_TOKEN]:
                    if not is_valid(token_list, next_token):
                        next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)

                    else:
                        token_list.append(next_token)
                        next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                # no consecutive whitespaces
                elif next_token == ' ' or next_token == '':
                    if whitespace == True:
                        next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)
                    else:
                        next_token, next_decoder_input_ids, decoder_input_ids, lm_logits, whitespace = add_and_gen_new(lm_logits, next_decoder_input_ids, decoder_input_ids, next_token, whitespace)
                # mask all other tokens
                else:
                    next_token, next_decoder_input_ids, lm_logits = mask_and_regen(lm_logits, next_decoder_input_ids)

            return [' '.join(token_list[:-1])]

    def parameters(self):
        return self.model.parameters()


# if __name__ == "__main__":
#     parser = argparse.ArgumentParser()
#     parser.add_argument("--data_fpath", type=str, default="data/holdout_split_batch12_perm/symbolic_batch12_perm_utt_0.2_0.pkl", help="complete file path or prefix of file paths to train test splits.")
#     parser.add_argument("--model_dpath", type=str, default=None, help="directory to save model checkpoints.")
#     parser.add_argument("--model_name", type=str, default="t5-large", choices=S2S_MODELS, help="name of supervised seq2seq model.")
#     parser.add_argument("--checkpoint", type=str, default=None, help="checkpoint to use for inferance.")
#     parser.add_argument("--results_dpath", type=str, default=None, help="directory to save evaluation results.")
#     args = parser.parse_args()

#     ckpt_dname = f"checkpoint-{args.checkpoint}" if args.checkpoint else "checkpoint-best"

#     logging.basicConfig(level=logging.DEBUG,
#                         format='%(message)s',
#                         handlers=[
#                             logging.FileHandler(f'{args.results_dpath}/s2s-tcd_{args.model_name}_{ckpt_dname}_{Path(args.data_fpath).stem}.log', mode='w'),
#                             logging.StreamHandler()
#                         ]
#     )

#     if "pkl" in args.data_fpath:  # complete file path, e.g. data/holdout_split_batch12_perm/symbolic_batch12_perm_utt_0.2_0.pkl
#         data_fpaths = [args.data_fpath]
#     else:  # prefix of file paths, e.g. for all utterance holdout datasets: data/holdout_split_batch12_perm/symbolic_batch12_perm_utt
#         data_dpath = os.path.dirname(args.data_fpath)
#         fname_prefix = os.path.basename(args.data_fpath)
#         data_fpaths = [os.path.join(data_dpath, fpath) for fpath in os.listdir(data_dpath) if fname_prefix in fpath]

#     for data_fpath in data_fpaths:
#         logging.info(f"Inference dataset: {data_fpath}")
#         # Load train, test data
#         train_iter, train_meta, valid_iter, valid_meta = load_split_dataset(data_fpath)

#         # Load trained model
#         if "t5" in args.model_name or "bart" in args.model_name:  # pretrained T5/Bart from Hugging Face
#             model_dpath = os.path.join(args.model_dpath, ckpt_dname)
#             logging.info(f"Load model and checkpoint: {model_dpath}")
#             s2s = Seq2Seq(model_dpath, args.model_name)
#         else:
#             raise TypeError(f"ERROR: unrecognized model, {args.model_name}")
#         logging.info(f"Number of trainable parameters in {args.model_name}: {count_params(s2s)}")
#         logging.info(f"Number of training samples: {len(train_iter)}")
#         logging.info(f"Number of validation samples: {len(valid_iter)}\n")

#         # Evaluate and save results
#         os.makedirs(args.results_dpath, exist_ok=True)
#         result_log_fpath = f"{args.results_dpath}/s2s-tcd_{args.model_name}_{ckpt_dname}_{Path(data_fpath).stem}_log.csv"
#         analysis_fpath = f"{os.path.expanduser('~')}/data/shared/lang2ltl/data/analysis_symbolic_batch12_perm.csv"
#         acc_fpath = f"{args.results_dpath}/s2s-tcd_{args.model_name}_{ckpt_dname}_{Path(data_fpath).stem}_acc.csv"
#         evaluate_sym_trans(s2s, data_fpath, result_log_fpath, analysis_fpath, acc_fpath, batch_size=50)