eval.py

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.distributions.categorical import Categorical
import os
import sys
import argparse
import numpy as np
import math
from copy import deepcopy
from model import BIML, describe_model
import datasets as dat
from train_lib import seed_all, extract, display_input_output, assert_consist_langs

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

## Evaluate a pre-trained model


def evaluate_ll(val_dataloader, net, langs, loss_fn=[], p_lapse=0.0, verbose=False):
    # Evaluate the total (sum) log-likelihood across the entire validation set
    #
    # Input
    #   val_dataloader :
    #   net : BIML model
    #   langs : dict of dat.Lang classes
    #   p_lapse : (default 0.) combine decoder outputs (prob 1-p_lapse) as mixture with uniform distribution (prob p_lapse)
    net.eval()
    total_N = 0
    total_ll = 0
    if not loss_fn:
        loss_fn = torch.nn.CrossEntropyLoss(ignore_index=langs["output"].PAD_idx)
    for batch_idx, val_batch in enumerate(val_dataloader):
        val_batch = dat.set_batch_to_device(val_batch)
        dict_loss = batch_ll(val_batch, net, loss_fn, langs, p_lapse=p_lapse)
        total_ll += dict_loss["ll"]
        total_N += dict_loss["N"]
    return total_ll, total_N


def evaluate_acc(
    val_dataloader, net, langs, max_length, eval_type="max", verbose=False
):
    # Evaluate accuracy (exact match) across entire validation set
    #
    # Input
    #   val_dataloader :
    #   net : BIML model
    #   langs : dict of dat.Lang classes
    #   max_length : maximum length of output sequences
    #   langs : dict of dat.Lang classes
    #   eval_type : 'max' for greedy decoding, 'sample' for sample from distribution
    #   out_mask_allow : default=[]; set of emission symbols we want to allow. Default of [] allows all output emissions
    net.eval()
    samples_pred = []  # list of all episodes with model predictions
    for batch_idx, val_batch in enumerate(val_dataloader):  # each batch
        val_batch = dat.set_batch_to_device(val_batch)
        scores = batch_acc(
            val_batch,
            net,
            langs,
            max_length,
            eval_type=eval_type,
            out_mask_allow=dat.get_batch_output_pool(val_batch),
        )
        samples_batch = val_batch["list_samples"]
        for sidx in range(len(samples_batch)):  # for each episode of the batch
            yq_sel = (
                val_batch["q_idx"].cpu().numpy() == sidx
            )  # select for queries in this episode
            in_support = scores["in_support"][yq_sel]  # numpy array
            is_novel = np.logical_not(in_support)
            v_acc = scores["v_acc"][yq_sel]  # numpy array
            samples_batch[sidx]["yq_predict"] = extract(yq_sel, scores["yq_predict"])
            samples_batch[sidx]["v_acc"] = v_acc
            samples_batch[sidx]["in_support"] = in_support  # numpy array
            samples_batch[sidx]["acc_retrieve"] = np.mean(v_acc[in_support]) * 100.0
            samples_batch[sidx]["acc_novel"] = np.mean(v_acc[is_novel]) * 100.0
        samples_pred.extend(samples_batch)

    # Compute mean accuracy across all val episodes
    mean_acc_retrieve = np.mean([sample["acc_retrieve"] for sample in samples_pred])
    v_acc_novel = [sample["acc_novel"] for sample in samples_pred]
    mean_acc_novel = np.mean(v_acc_novel)

    if verbose:
        display_console_pred(samples_pred)
    return {
        "samples_pred": samples_pred,
        "mean_acc_novel": mean_acc_novel,
        "mean_acc_retrieve": mean_acc_retrieve,
        "v_novel": v_acc_novel,
    }


def batch_ll(batch, net, loss_fn, langs, p_lapse=0.0):
    # Evaluate log-likelihood (average over cells, and sum total) for a given batch
    #
    # Input
    #   batch : from dat.make_biml_batch
    #   loss_fn : loss function
    #   langs : dict of dat.Lang classes
    net.eval()
    m = len(batch["yq"])  # b*nq
    target_batches = batch["yq_padded"]  # b*nq x max_length
    target_lengths = batch["yq_lengths"]  # list of size b*nq
    target_shift = batch["yq_sos_padded"]  # b*nq x max_length
    # Shifted targets with padding (added SOS symbol at beginning and removed EOS symbol)
    decoder_output = net(target_shift, batch)
    # b*nq x max_length x output_size

    logits_flat = decoder_output.reshape(
        -1, decoder_output.shape[-1]
    )  # (batch*max_len, output_size)
    if p_lapse > 0:
        logits_flat = smooth_decoder_outputs(
            logits_flat, p_lapse, langs["output"].symbols + [dat.EOS_token], langs
        )
    loss = loss_fn(logits_flat, target_batches.reshape(-1))
    loglike = -loss.cpu().item()
    dict_loss = {}
    dict_loss["ll_by_cell"] = loglike  # average over cells
    dict_loss["N"] = float(sum(target_lengths))  # total number of valid cells
    dict_loss["ll"] = dict_loss["ll_by_cell"] * dict_loss["N"]  # total LL
    return dict_loss


def smooth_decoder_outputs(logits_flat, p_lapse, lapse_symb_include, langs):
    # Mix decoder outputs (logits_flat) with uniform distribution over allowed emissions (in lapse_symb_include)
    #
    # Input
    #  logits_flat : (batch*max_len, output_size) # unnomralized log-probabilities
    #  p_lapse : probability of a uniform lapse
    #  lapse_symb_include : list of tokens (strings) that we want to include in the lapse model
    #  langs : dict of dat.Lang classes
    #
    # Output
    #  log_probs_flat : (batch*max_len, output_size) normalized log-probabilities
    lapse_idx_include = [langs["output"].symbol2index[s] for s in lapse_symb_include]
    assert (
        dat.SOS_token not in lapse_symb_include
    )  # SOS should not be an allowed output through lapse model
    sz = logits_flat.size()  # get size (batch*max_len, output_size)
    probs_flat = F.softmax(logits_flat, dim=1)  # (batch*max_len, output_size)
    num_classes_lapse = len(lapse_idx_include)
    probs_lapse = torch.zeros(sz, dtype=torch.float)
    probs_lapse = probs_lapse.to(device=DEVICE)
    probs_lapse[:, lapse_idx_include] = 1.0 / float(num_classes_lapse)
    log_probs_flat = torch.log(
        (1 - p_lapse) * probs_flat + p_lapse * probs_lapse
    )  # (batch*max_len, output_size)
    return log_probs_flat


def batch_acc(batch, net, langs, max_length, eval_type="max", out_mask_allow=[]):
    # Evaluate exact match accuracy for a given batch
    #
    #  Input
    #   batch : from dat.make_biml_batch
    #   net : BIML model
    #   max_length : maximum length of output sequences
    #   langs : dict of dat.Lang classes
    #   eval_type : 'max' for greedy decoding, 'sample' for sample from distribution
    #   out_mask_allow : default=[]; list of emission symbols (strings) we want to allow. Default of [] allows all output emissions
    assert eval_type in ["max", "sample"]
    net.eval()
    emission_lang = langs["output"]
    use_mask = len(out_mask_allow) > 0
    memory, memory_padding_mask = net.encode(batch)
    # memory : b*nq x maxlength_src x hidden_size
    # memory_padding_mask : b*nq x maxlength_src (False means leave alone)
    m = len(batch["yq"])  # b*nq
    z_padded = torch.tensor(
        [emission_lang.symbol2index[dat.SOS_token]] * m
    )  # b*nq length tensor
    z_padded = z_padded.unsqueeze(1)  # [b*nq x 1] tensor
    z_padded = z_padded.to(device=DEVICE)
    max_length_target = batch["yq_padded"].shape[1] - 1  # length without EOS
    assert (
        max_length >= max_length_target
    )  # make sure that the net can generate targets of the proper length

    # make the output mask if certain emissions are restricted
    if use_mask:
        assert (
            dat.EOS_token in out_mask_allow
        )  # EOS must be included as an allowed symbol
        additive_out_mask = -torch.inf * torch.ones(
            (m, net.output_size), dtype=torch.float
        )
        additive_out_mask = additive_out_mask.to(device=DEVICE)
        for s in out_mask_allow:
            sidx = langs["output"].symbol2index[s]
            additive_out_mask[:, sidx] = 0.0

    # Run through decoder
    all_decoder_outputs = torch.zeros((m, max_length), dtype=torch.long)
    all_decoder_outputs = all_decoder_outputs.to(device=DEVICE)
    for t in range(max_length):
        decoder_output = net.decode(z_padded, memory, memory_padding_mask)
        # decoder_output is b*nq x (t+1) x output_size
        decoder_output = decoder_output[
            :, -1
        ]  # get the last step's output (batch_size x output_size)
        if use_mask:
            decoder_output += additive_out_mask

        # Choose the symbols at next timestep
        if eval_type == "max":  # pick the most likely
            topi = torch.argmax(decoder_output, dim=1)
            emissions = topi.view(-1)
        elif eval_type == "sample":
            emissions = Categorical(logits=decoder_output).sample()
        all_decoder_outputs[:, t] = emissions
        z_padded = torch.cat([z_padded, emissions.unsqueeze(1)], dim=1)

    # Get predictions as strings and see if they are correct
    all_decoder_outputs = all_decoder_outputs.detach()
    yq_predict = []  # list of all predicted query outputs as strings
    v_acc = np.zeros(m)
    for q in range(m):
        myseq = emission_lang.tensor_to_symbols(all_decoder_outputs[q, :].view(-1))
        yq_predict.append(myseq)
        v_acc[q] = (
            yq_predict[q] == batch["yq"][q]
        )  # for each query, did model get it right?
    in_support = np.array(batch["in_support"])  # which queries are also support items
    out = {"yq_predict": yq_predict, "v_acc": v_acc, "in_support": in_support}
    return out


def viz_train_dashboard(train_tracker):
    # Show loss curves
    import matplotlib.pyplot as plt

    if not train_tracker:
        print("No training stats to plot")
        return
    fv = lambda x: [t[x] for t in train_tracker]
    plt.figure(1)
    plt.subplot(2, 2, 1)
    plt.plot(fv("step"), fv("avg_train_loss"), "b", label="train")
    if "val_loss" in train_tracker[0]:
        plt.plot(fv("step"), fv("val_loss"), "r", label="val")
    plt.xlabel("step")
    plt.legend()
    plt.title("Loss")
    plt.subplot(2, 2, 2)
    plt.plot(fv("step"), fv("lr"), "b")
    plt.xlabel("step")
    plt.title("Learning rate")
    plt.show()


def display_console_pred(samples_pred):
    # Print model predictions
    #
    # Input
    #  samples_pred : list of dicts from evaluate_acc, which has predicted query outputs for each episode
    for idx, sample in enumerate(samples_pred):
        print("Evaluation episode " + str(idx))
        in_support = sample["in_support"]
        is_novel = np.logical_not(in_support)
        if "grammar" in sample:
            print("")
            print(sample["grammar"])
        print("  support items;")
        display_input_output(sample["xs"], sample["ys"], sample["ys"])
        print("  retrieval items;", round(sample["acc_retrieve"], 3), "% correct")
        display_input_output(
            extract(in_support, sample["xq"]),
            extract(in_support, sample["yq_predict"]),
            extract(in_support, sample["yq"]),
        )
        print("  generalization items;", round(sample["acc_novel"], 3), "% correct")
        display_input_output(
            extract(is_novel, sample["xq"]),
            extract(is_novel, sample["yq_predict"]),
            extract(is_novel, sample["yq"]),
        )


def display_console_unmap(samples_pred):
    # Print model predictions after remapping
    #  There must also be a remapping from current tokens back to canonical tokens/
    #
    # Input
    #  samples_pred : list of dicts from evaluate_acc, which has predictions for each episode
    for idx, sample in enumerate(samples_pred):
        assert (
            "unmap_input" in sample["aux"]
        ), "there must be mapping back to canonical text form"
        ui = lambda x: list(map(sample["aux"]["unmap_input"], x))
        uo = lambda x: list(map(sample["aux"]["unmap_output"], x))
        if "filename" in sample["aux"]:
            print(
                "Evaluation episode " + str(idx) + "; filename:",
                sample["aux"]["filename"],
            )
        else:
            print("Evaluation episode " + str(idx))
        in_support = sample["in_support"]
        is_novel = np.logical_not(in_support)
        if "grammar" in sample:
            print("")
            print(sample["grammar"])
        print("  support items;")
        display_input_output(ui(sample["xs"]), uo(sample["ys"]), uo(sample["ys"]))
        print("  retrieval items;", round(sample["acc_retrieve"], 3), "% correct")
        display_input_output(
            extract(in_support, ui(sample["xq"])),
            extract(in_support, uo(sample["yq_predict"])),
            extract(in_support, uo(sample["yq"])),
        )
        print("  generalization items;", round(sample["acc_novel"], 3), "% correct")
        display_input_output(
            extract(is_novel, ui(sample["xq"])),
            extract(is_novel, uo(sample["yq_predict"])),
            extract(is_novel, uo(sample["yq"])),
        )


def display_html_unmap(samples_pred, fid, freq="percent", include_support=False):
    #  Show model predictions when sampling.
    #   Each episode consists of just one command, repeated multiple times.
    #   There must also be a remapping back to canonical text format.
    #
    #   Input:
    #     samples_pred : list of dicts from evaluate_acc, which has predictions for each episode
    #     fid : handle for text file we are writing to
    #     freq : [percent OR count], format for reporting frequency
    #     include_support : show support set in HTML format? used for probe task
    fid.write("var all_data = [")
    for idx, sample in enumerate(samples_pred):
        assert freq in ["percent", "count"]
        assert (
            "unmap_input" in sample["aux"]
        ), "there must be mapping back to canonical text form"
        ui = lambda x: list(map(sample["aux"]["unmap_input"], x))
        uo = lambda x: list(map(sample["aux"]["unmap_output"], x))
        in_support = sample["in_support"]
        xq_novel = extract(np.logical_not(in_support), ui(sample["xq"]))
        yq_predict_novel = extract(np.logical_not(in_support), uo(sample["yq_predict"]))
        xq_novel = [" ".join(x) for x in xq_novel]
        assert all(
            xq == xq_novel[0] for xq in xq_novel
        ), "each episode must be repeats of the same command"
        mycommand = xq_novel[0]
        myresponses = [" ".join(y) for y in yq_predict_novel]
        unique_responses = sorted(set(myresponses))
        count_unique = []
        for u in unique_responses:
            if freq == "percent":
                count_unique.append(
                    100.0
                    * np.mean(np.array([u == rr for rr in myresponses], dtype=float))
                )
            else:
                count_unique.append(
                    np.sum(np.array([u == rr for rr in myresponses], dtype=int))
                )
        unique_responses = [
            x
            for x in sorted(
                zip(unique_responses, count_unique),
                key=lambda pair: pair[1],
                reverse=True,
            )
        ]
        if include_support:
            fid.write("[ \n")
            fid.write("'support', \n")
            xs = ui(sample["xs"])
            ys = uo(sample["ys"])
            for j in range(len(xs)):
                fid.write("  ['%s', '%s'], \n" % (" ".join(xs[j]), " ".join(ys[j])))
            if "output_pool" in sample["aux"]:
                mypool = sample["aux"]["output_pool"]
                fid.write("  ['%s', '%s'], \n" % ("Pool", " ".join(uo([mypool])[0])))
            fid.write("], \n")
        fid.write("[ \n")
        if include_support:
            fid.write("'query', \n")
        for u in unique_responses:
            if freq == "percent":
                fid.write("  ['%s', '%s', '%.1f%%'], \n" % (mycommand, u[0], u[1]))
            else:
                fid.write("  ['%s', '%s', %d], \n" % (mycommand, u[0], u[1]))
        fid.write("], \n")
    fid.write("]; \n")


def evaluate_ll_multi_rep(nrep, val_dataloader, net, langs, p_lapse=0.0):
    # Across multiple random runs, evaluate the total (sum) log-likelihood across the validation set.
    #   Multiple runs may be needed when certain things are randomized, like the mapping between original data files
    #   and input/output tokens.
    #
    #  Input
    #   nrep : number of replications
    #    for rest of input arguments, see def evaluate_ll
    #   net : BIML model
    #   langs : dict of dat.Lang classes
    #
    # Output
    #   ave_ll_by_cell : return average log-likelihood for each cell
    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=langs["output"].PAD_idx)
    list_ll = []
    list_N = []
    for i in range(nrep):
        if nrep >= 5 and (i % int(nrep / 5) == 0):
            print("  run", i)
        total_ll, total_N = evaluate_ll(
            val_dataloader, net, langs, loss_fn=loss_fn, p_lapse=p_lapse
        )
        list_ll.append(total_ll)
        list_N.append(total_N)
    mean_ll = np.mean(list_ll)
    std_ll = np.std(list_ll)
    mean_N = np.mean(list_N)
    std_N = np.std(list_N)
    ave_ll_by_cell = np.sum(list_ll) / float(np.sum(list_N))
    print(
        "  loglike: M =",
        round(mean_ll, 4),
        "(SD=",
        round(std_ll, 4),
        ", Nrep=",
        nrep,
        ") for",
        round(mean_N, 4),
        "(SD=",
        round(std_N, 4),
        ") symbol predictions on average",
    )
    print("    ave LL by cell: ", round(ave_ll_by_cell, 4))
    return {
        "ave_ll_by_cell": ave_ll_by_cell,
        "ave_ll": mean_ll,
        "std_ll": std_ll,
        "ave_N": mean_N,
    }


def fit_p_lapse(nrep, val_dataloader, net, langs, greedy_stop=True):
    # Fit value for p_lapse. For each value, use evaluate_ll_multi_rep
    #  greedy_stop : stop when the objective no longer improves
    #
    # Output
    #  p_lapse : best fitting lapse value
    #  score_best : average log-like of best value
    #  mean_N : average number of evaluations
    iter_p_lapse = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09]
    iter_p_lapse += [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    score_best = {"ave_ll_by_cell": float(-math.inf)}
    for p_lapse in iter_p_lapse:
        print(" p_lapse", p_lapse, ":")
        seed_all()
        score_curr = evaluate_ll_multi_rep(
            nrep, val_dataloader, net, langs, p_lapse=p_lapse
        )
        if score_curr["ave_ll_by_cell"] >= score_best["ave_ll_by_cell"]:
            score_best = score_curr
            p_lapse_best = p_lapse
        elif greedy_stop:
            break  #
    return p_lapse_best, score_best


def evaluate_iterative(
    val_dataloader,
    net,
    langs,
    max_length,
    eval_type="max",
    out_mask_allow=[],
    verbose=False,
):
    # Sample from model iteratively:
    #   1) Generate output for the first query
    #   2) Add this query as a study example, using self-generated output as the target
    #   3) If there is another query, go back to step 1.
    #
    # Input
    #   val_dataloader : PyTorch dataloader
    #   net : BIML model
    #   langs : dict of dat.Lang classes
    #   max_length : maximum length of the output sequences
    #   eval_type : 'max' for greedy decoding, 'sample' for sample from distribution
    #   out_mask_allow : default=[]; set of emission symbols we want to allow. Default of [] allows all output emissions
    net.eval()
    samples_pred = []  # list of all episodes with model predictions
    for batch_idx, val_batch in enumerate(val_dataloader):
        val_batch = dat.set_batch_to_device(val_batch)
        scores, val_batch = batch_iterative(
            val_batch,
            net,
            langs,
            max_length,
            eval_type=eval_type,
            out_mask_allow=out_mask_allow,
            verbose=verbose,
        )
        samples_batch = val_batch["list_samples"]
        for sidx in range(len(samples_batch)):  # for each episode of the batch
            yq_sel = (
                val_batch["q_idx"].cpu().numpy() == sidx
            )  # select for queries in this episode
            m = np.sum(yq_sel)
            samples_batch[sidx]["yq_predict"] = extract(yq_sel, scores["yq_predict"])
            samples_batch[sidx]["in_support"] = np.ones(m, dtype=bool)
            samples_batch[sidx]["acc_retrieve"] = float("nan")
            samples_batch[sidx]["acc_novel"] = float("nan")
        samples_pred.extend(samples_batch)
    return {"samples_pred": samples_pred}


def batch_iterative(
    batch, net, langs, max_length, eval_type="max", out_mask_allow=[], verbose=False
):
    # Helper function for evaluate_iterative. Processes a whole batch in iterative manner
    batch_next = batch
    flag_batch_changed = True
    ii = 0
    while flag_batch_changed:
        # Model predicts outputs for each query
        scores = batch_acc(
            batch_next,
            net,
            langs,
            max_length,
            eval_type=eval_type,
            out_mask_allow=out_mask_allow,
        )
        samples_curr = batch_next["list_samples"]  # list of samples
        samples_next = samples_curr  # placeholder for modified samples
        bsize = len(samples_curr)  # number of episodes
        for sidx in range(bsize):  # divide predictions by episodes
            yq_sel = (
                batch_next["q_idx"].cpu().numpy() == sidx
            )  # select for queries in this episode
            samples_curr[sidx]["yq_predict"] = extract(yq_sel, scores["yq_predict"])
            samples_curr[sidx]["in_support"] = scores["in_support"][
                yq_sel
            ]  # numpy array

        if verbose:
            print("\n** Iteration ", ii, "**")
            samples_pred = samples_curr
            for idx, sample in enumerate(samples_pred[:1]):
                print("Evaluation episode " + str(idx))
                in_support = sample["in_support"]
                print("  support items: ")
                display_input_output(sample["xs"], sample["ys"], sample["ys"])
                print("  retrieval items; ")
                display_input_output(
                    extract(in_support, sample["xq"]),
                    extract(in_support, sample["yq_predict"]),
                    extract(in_support, sample["yq"]),
                )
                print("  generalization items; ")
                display_input_output(
                    extract(np.logical_not(in_support), sample["xq"]),
                    extract(np.logical_not(in_support), sample["yq_predict"]),
                    extract(np.logical_not(in_support), sample["yq"]),
                )

        # Add command from first query to support set, using self-generated output as target
        flag_batch_changed = False  # did any episode get modified this iteration?
        samples_new_list = []
        for sidx in range(bsize):  # for each episode
            S_curr = deepcopy(samples_curr[sidx])  # current episode
            is_novel = np.logical_not(
                S_curr["in_support"]
            )  # queries that are not yet in the support set
            if np.any(is_novel):  # if there is a genuine query
                flag_batch_changed = True
                myidx = np.nonzero(is_novel)[0][0]  # pick the first query
                xq_add = S_curr["xq"][myidx]
                yq_add = S_curr["yq_predict"][myidx]

                # add query to support set
                samples_next[sidx]["xs"].append(deepcopy(xq_add))
                samples_next[sidx]["ys"].append(deepcopy(yq_add))

                # keep query in query set, but update the target
                del samples_next[sidx]["xq"][myidx]
                del samples_next[sidx]["yq"][myidx]
                samples_next[sidx]["xq"].append(deepcopy(xq_add))
                samples_next[sidx]["yq"].append(deepcopy(yq_add))
            samples_new_list += [
                dat.bundle_biml_episode(
                    samples_next[sidx]["xs"],
                    samples_next[sidx]["ys"],
                    samples_next[sidx]["xq"],
                    samples_next[sidx]["yq"],
                    "",
                )
            ]

        # create the next batch
        batch_next = dat.make_biml_batch(samples_new_list, langs)
        batch_next = dat.set_batch_to_device(batch_next)
        ii += 1
    return scores, batch_next


if __name__ == "__main__":
    # Adjustable parameters
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--fn_out_model",
        type=str,
        default="",
        help="*REQUIRED*. Filename for loading the model",
    )
    parser.add_argument(
        "--dir_model",
        type=str,
        default="out_models",
        help="Directory for loading the model file",
    )
    parser.add_argument(
        "--max_length_eval",
        type=int,
        default=10,
        help="Maximum generated sequence length",
    )
    parser.add_argument(
        "--batch_size", type=int, default=-1, help="Number of episodes in batch"
    )
    parser.add_argument(
        "--episode_type",
        type=str,
        default="",
        help="What type of episodes do we want? See datasets.py for options",
    )
    parser.add_argument(
        "--dashboard",
        default=False,
        action="store_true",
        help="Showing loss curves during training.",
    )
    parser.add_argument(
        "--ll",
        default=False,
        action="store_true",
        help="Evaluate log-likelihood of validation (val) set",
    )
    parser.add_argument(
        "--max",
        default=False,
        action="store_true",
        help="Find best outputs for val commands (greedy decoding)",
    )
    parser.add_argument(
        "--sample",
        default=False,
        action="store_true",
        help="Sample outputs for val commands",
    )
    parser.add_argument(
        "--sample_html",
        default=False,
        action="store_true",
        help="Sample outputs for val commands in html format (using unmap to canonical text)",
    )
    parser.add_argument(
        "--sample_iterative",
        default=False,
        action="store_true",
        help="Sample outputs for val commands iteratively. Output in html format",
    )
    parser.add_argument(
        "--fit_lapse",
        default=False,
        action="store_true",
        help="Fit the best lapse rate according to log-likelihood on validation",
    )
    parser.add_argument(
        "--ll_nrep",
        type=int,
        default=1,
        help="Evaluate each episode this many times when computing log-likelihood (needed for stochastic remappings)",
    )
    parser.add_argument(
        "--ll_p_lapse",
        type=float,
        default=0.0,
        help="Lapse rate when evaluating log-likelihoods",
    )
    parser.add_argument(
        "--verbose",
        default=False,
        action="store_true",
        help="Inspect outputs in more detail",
    )

    args = parser.parse_args()
    fn_out_model = args.fn_out_model
    dir_model = args.dir_model
    max_length_eval = args.max_length_eval
    episode_type = args.episode_type
    do_dashboard = args.dashboard
    batch_size = args.batch_size
    do_ll = args.ll
    do_max_acc = args.max
    do_sample_acc = args.sample
    do_sample_html = args.sample_html
    do_sample_iterative = args.sample_iterative
    do_fit_lapse = args.fit_lapse
    ll_nrep = args.ll_nrep
    ll_p_lapse = args.ll_p_lapse
    verbose = args.verbose

    model_tag = episode_type + "_" + fn_out_model.replace(".pt", "")
    fn_out_model = os.path.join(dir_model, fn_out_model)
    if not os.path.isfile(fn_out_model):
        raise Exception("filename " + fn_out_model + " not found")

    seed_all()
    print("Loading model:", fn_out_model, "on", DEVICE)
    checkpoint = torch.load(fn_out_model, map_location=DEVICE, weights_only=False)
    if not episode_type:
        episode_type = checkpoint["episode_type"]
    if batch_size <= 0:
        batch_size = checkpoint["batch_size"]
    nets_state_dict = checkpoint["nets_state_dict"]
    if list(nets_state_dict.keys()) == ["net"]:
        nets_state_dict = nets_state_dict["net"]  # for compatibility with legacy code
    input_size = checkpoint["langs"]["input"].n_symbols
    output_size = checkpoint["langs"]["output"].n_symbols
    emb_size = checkpoint["emb_size"]
    dropout_p = checkpoint["dropout"]
    ff_mult = checkpoint["ff_mult"]
    myact = checkpoint["activation"]
    nlayers_encoder = checkpoint["nlayers_encoder"]
    nlayers_decoder = checkpoint["nlayers_decoder"]
    train_tracker = checkpoint["train_tracker"]
    best_val_loss = -float("inf")
    if "best_val_loss" in checkpoint:
        best_val_loss = checkpoint["best_val_loss"]

    print(
        " Loading model that has completed (or started) "
        + str(checkpoint["epoch"])
        + " of "
        + str(checkpoint["nepochs"])
        + " epochs"
    )
    print("  test episode_type:", episode_type)
    print("  batch size:", checkpoint["batch_size"])
    print("  max eval length:", max_length_eval)
    print("  number of steps:", checkpoint["step"])
    print("  best val loss achieved: {:.4f}".format(best_val_loss))

    # Load validation dataset
    D_train, D_val = dat.get_dataset(episode_type)
    langs = D_val.langs
    assert_consist_langs(langs, checkpoint["langs"])
    train_dataloader = DataLoader(
        D_train,
        batch_size=batch_size,
        collate_fn=lambda x: dat.make_biml_batch(x, langs),
        shuffle=False,
    )
    val_dataloader = DataLoader(
        D_val,
        batch_size=batch_size,
        collate_fn=lambda x: dat.make_biml_batch(x, langs),
        shuffle=False,
    )

    # For backward compatibility with legacy code that used same EOS and PAD tokens
    add_pad = dat.PAD_token not in checkpoint["langs"]["input"].symbol2index

    # Load model parameters
    net = BIML(
        emb_size,
        input_size,
        output_size,
        langs["input"].PAD_idx,
        langs["output"].PAD_idx,
        nlayers_encoder=nlayers_encoder,
        nlayers_decoder=nlayers_decoder,
        dropout_p=dropout_p,
        activation=myact,
        ff_mult=ff_mult,
    )
    net.load_state_dict(nets_state_dict)
    net = net.to(device=DEVICE)
    describe_model(net)

    # Perform selected evaluations
    if do_dashboard:
        print("Showing loss curves during training <close plot to continue>")
        viz_train_dashboard(train_tracker)
    if do_ll and ll_nrep == 1:
        seed_all()
        print("Evaluating log-likelihood of val episodes...")
        print("  with lapse rate", ll_p_lapse)
        total_ll, total_N = evaluate_ll(val_dataloader, net, langs, p_lapse=ll_p_lapse)
        print(
            "evaluation on",
            episode_type,
            "loglike:",
            round(total_ll, 4),
            "for",
            int(total_N),
            "symbol predictions",
        )
        print("mean loglike is", round(total_ll / total_N, 5), "per symbol")
    if do_ll and ll_nrep > 1:
        seed_all()
        print("Evaluating log-likelihood of val episodes...")
        print("  with lapse rate", ll_p_lapse)
        print("  replicated across", ll_nrep, "random runs/permutations")
        evaluate_ll_multi_rep(ll_nrep, val_dataloader, net, langs, p_lapse=ll_p_lapse)
    if do_max_acc:
        seed_all()
        E = evaluate_acc(
            val_dataloader,
            net,
            langs,
            max_length_eval,
            eval_type="max",
            verbose=verbose,
        )
        print("Evaluating set of validation episodes (via greedy decoding)...")
        print(" Acc Retrieve (val):", round(E["mean_acc_retrieve"], 4))
        print(
            " Acc Novel (val):",
            round(np.mean(E["v_novel"]), 4),
            "SD=",
            round(np.std(E["v_novel"]), 4),
            "N=",
            len(E["v_novel"]),
        )
    if do_sample_acc:
        seed_all()
        E = evaluate_acc(
            val_dataloader,
            net,
            langs,
            max_length_eval,
            eval_type="sample",
            verbose=verbose,
        )
        print("Evaluating set of validation episodes (via sampling)...")
        print(" Acc Retrieve (val):", round(E["mean_acc_retrieve"], 4))
        print(
            " Acc Novel (val):",
            round(np.mean(E["v_novel"]), 4),
            "SD=",
            round(np.std(E["v_novel"]), 4),
            "N=",
            len(E["v_novel"]),
        )
    if do_sample_html:
        assert episode_type in ["few_shot_human_mult10"]
        seed_all()
        print("Sampling from model to produce HTML file...")
        E = evaluate_acc(
            val_dataloader,
            net,
            langs,
            max_length_eval,
            eval_type="sample",
            verbose=False,
        )
        episode_type_tag = episode_type
        if "probe_human" in episode_type:
            episode_type_tag = "probe_human"
        with open(
            "html_output/" + episode_type_tag + "/" + model_tag + ".txt", "w"
        ) as fid_out:
            prev = sys.stdout
            sys.stdout = fid_out  # re-rout outputs to file
            display_console_unmap(E["samples_pred"])
            sys.stdout = prev
        with open("html_output/" + episode_type_tag + "/template.html", "r") as fid_in:
            mylines = fid_in.readlines()
        with open(
            "html_output/" + episode_type_tag + "/" + model_tag + ".html", "w"
        ) as fid_out:
            for l in mylines:
                fid_out.write(l)
                if l.strip() == "// PLACEHOLDER":
                    fid_out.write('var title="' + model_tag + '"; \n')
                    if "probe" in episode_type_tag:
                        display_html_unmap(
                            E["samples_pred"],
                            fid_out,
                            freq="percent",
                            include_support=True,
                        )
                    else:
                        display_html_unmap(E["samples_pred"], fid_out, freq="percent")
        print(
            "Done writing to files: html_output/"
            + episode_type_tag
            + "/"
            + model_tag
            + ".html and .txt"
        )
    if do_sample_iterative:
        assert episode_type in ["open_end_freeform"]
        seed_all()
        print(
            "Iteratively evaluate queries by adding them, one-by-one, to the support set using self-generated targets."
        )
        E = evaluate_iterative(
            val_dataloader,
            net,
            langs,
            max_length_eval,
            eval_type="sample",
            verbose=verbose,
        )
        with open(
            "html_output/" + episode_type + "/" + model_tag + ".txt", "w"
        ) as fid_out:
            prev = sys.stdout
            sys.stdout = fid_out  # re-rout outputs to file
            display_console_pred(E["samples_pred"])
            sys.stdout = prev
    if do_fit_lapse:
        print("Fitting for the best value of p_lapse use log-like...")
        print("  Each value is replicated across", ll_nrep, "random runs/permutations")
        assert ll_nrep > 1, "should use more than one replication for parameter fitting"
        p_lapse_best, score_best = fit_p_lapse(ll_nrep, val_dataloader, net, langs)
        print(
            "* BEST FIT * p_lapse=",
            p_lapse_best,
            "with mean loglike score of",
            round(score_best["ave_ll"], 4),
            "(or",
            round(score_best["ave_ll_by_cell"], 4),
            "per cell)",
        )