text:
# python reorder_single.py --bs $BS --alpha $alpha --insert

import argparse
import copy
import gc
import itertools
import json
import random
import sys
from itertools import permutations
from os import sep

import numpy as np
import pandas as pd
import torch
import torch as T
from torch.autograd import grad

import wandb
from dlg_utils import calculateGradLoss, calculatePerplexity, get_grad_gpt2
from eval_utils import eval

stride = 16

def shard_sequence_using_eos(input, sep_token_id=50256):
    ret = []
    start = 0
    for i in range(len(input)):
        if input[i] == sep_token_id:
            ret.append(input[start:i])
            start = i + 1
    ret.append(input[start:])
    return ret

def score(
    model, padded, padded_s, original_dy_dx, tokenizer_GPT, length=20, batchsize=1
):
    padded = padded.long()

    gradscore = []

    for idx, s in enumerate(padded_s):
        if batchsize == 1:
            s = s.replace("<|endoftext|>", "")
        else:
            p = shard_sequence_using_eos(padded[idx][1:-1])
            s = tokenizer_GPT.batch_decode(p)  # Seperate into sentences of size
        padded_dy_dx = get_grad_gpt2(s, tokenizer_GPT, model)
        gradnorm = torch.norm(padded_dy_dx[LAYER_IDX]).cpu().numpy().item()
        ppl = calculatePerplexity(s, model, tokenizer_GPT).item()

        gradscore.append(gradnorm + ALPHA * ppl)
    gradscore = np.asarray(gradscore)

    wordscoress, scoress, predss = [], [], []
    for i in range(int(np.ceil(len(padded) / stride))):
        outputs = model(padded[i * stride : min((i + 1) * stride, len(padded))])
        lsm = -outputs[0].log_softmax(2)
        preds = T.zeros_like(lsm)
        preds[:, 1:] = lsm[:, :-1]
        wordscores = (
            preds.gather(
                2, padded[i * stride : min((i + 1) * stride, len(padded))].unsqueeze(2)
            )
            .squeeze(2)
            .cpu()
            .detach()
        )
        scores = wordscores.sum(1)

        wordscoress.append(wordscores)
        scoress.append(scores)
        predss.append(preds.cpu().detach())

    wordscores = T.cat(wordscoress)
    scores = T.cat(scoress)
    preds = T.cat(predss)

    return gradscore, scores, wordscores, -preds

cand_orders = {
    3: [
        list(ll)
        for ll in list(permutations(range(1, 5)))
        if ll[0] == 1 and ll[-1] == 4 and ll != tuple(range(1, 5))
    ],
    4: [
        list(ll)
        for ll in list(permutations(range(1, 6)))
        if ll[0] == 1 and ll[-1] == 5 and ll != tuple(range(1, 6))
    ],
    5: [
        list(ll)
        for ll in list(permutations(range(1, 7)))
        if ll[0] == 1 and ll[-1] == 6 and ll != tuple(range(1, 7))
    ],
    6: [
        list(ll)
        for ll in list(permutations(range(1, 8)))
        if ll[0] == 1 and ll[-1] == 6 and ll != tuple(range(1, 8))
    ],
    7: [
        list(ll)
        for ll in list(permutations(range(1, 9)))
        if ll[0] == 1 and ll[-1] == 6 and ll != tuple(range(1, 9))
    ],
}

class LocalSwap:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def computeStats(self, curr_tokens, original_dy_dx, do_split=False):

        generated_sentence = self.tokenizer.decode(curr_tokens)
        if do_split:
            generated_sentence = generated_sentence.split("<|endoftext|>")

        generated_dy_dx = get_grad_gpt2(generated_sentence, self.tokenizer, self.model)

        gradnorm = torch.norm(generated_dy_dx[LAYER_IDX]).cpu().numpy().item()
        ppl = calculatePerplexity(
            generated_sentence[0], self.model, self.tokenizer
        ).item()
        return {
            "tokens": curr_tokens,
            "sentence": generated_sentence,
            "gradloss_mse": gradnorm + ALPHA * ppl,
        }

    def createAllOrderings(
        self, sentence_size=40, without_order=False, max_depth=5, sep_idx=None
    ):
        for d in range(1, max_depth):
            tuples = itertools.product(range(1, sentence_size), range(1, sentence_size))
            if sep_idx == []:
                tuples = (
                    t for t in tuples if t[0] != t[1]
                )  # Only swap inside the first sentence! IMPORTANT!
                # if t[0] != t[1] and (t[0] not in sep_idx) and (t[1] not in sep_idx)
            else:
                tuples = (
                    t
                    for t in tuples
                    if t[0] != t[1]
                    and (t[0] < sep_idx[0])
                    and (
                        t[1] < sep_idx[0]
                    )  # Only swap inside the first sentence! IMPORTANT!
                    # if t[0] != t[1] and (t[0] not in sep_idx) and (t[1] not in sep_idx)
                )
            if without_order:
                tuples = (t for t in tuples if t[0] < t[1])
            tuples = list(tuples)

            tuple_list = [copy.copy(tuples) for _ in range(d)]

            for t in tuple_list:
                random.shuffle(t)

            for ordering in itertools.product(*tuple_list):
                yield ordering

    def reorderAlgorithm(self, original_dy_dx, tokens, generator, n=50):
        values = []
        for _, sequence in zip(range(n), generator):
            curr_tokens = copy.copy(tokens)
            for (i, j) in sequence:
                curr_tokens[i], curr_tokens[j] = curr_tokens[j], curr_tokens[i]
            values.append(self.computeStats(curr_tokens, original_dy_dx, do_split=True))

        if len(values) == 0:
            return {"gradloss_mse": 10e100}

        return min(values, key=lambda x: x["gradloss_mse"])

    def insertAlgorithm(
        self, original_dy_dx, tokens, token_pool, generator, n=50
    ):  # insert the i-th token from the pool at position
        if generator == None:  # insertion disabled
            return {"gradloss_mse": 10e100}
        if len(tokens) >= 40:  # disable insertion if the sentence is alreday very long
            return {"gradloss_mse": 10e100}
        values = []
        for _, sequence in zip(range(n), generator):
            curr_tokens = copy.copy(tokens)

            for (i, j) in sequence:
                # if i < j:
                #     curr_tokens = (
                #         curr_tokens[0:i]
                #         + curr_tokens[i + 1 : j + 1]
                #         + [curr_tokens[i]]
                #         + curr_tokens[j + 1 :]
                #     )
                # else:
                if i < len(token_pool) and j < len(tokens):
                    curr_tokens = curr_tokens[0:j] + [token_pool[i]] + curr_tokens[j:]
            values.append(self.computeStats(curr_tokens, original_dy_dx, do_split=True))

        if len(values) == 0:
            return {"gradloss_mse": 10e100}

        return min(values, key=lambda x: x["gradloss_mse"])

def shuffle_proposals_rand(length, bs, kopt):
    L = length
    orders = []
    o = np.arange(2, kopt + 1)
    np.random.shuffle(o)
    o = np.insert(o, 0, 1)
    o = np.insert(o, len(o), kopt + 1)
    orders = [o for _ in range(bs)]

    imod = [(torch.randperm(L - 2) + 1)[:kopt].sort()[0] for _ in range(bs)]
    topv = torch.zeros(bs)

    return T.stack(imod), topv, orders

def shuffle_proposals(mat, topk, bs, kopt):

    L = mat.shape[0]

    I = T.zeros((kopt,) + (L,) * (kopt)).long()
    for i in range(kopt):
        I[i] = T.arange(L).view((-1,) + (1,) * (kopt - 1 - i))

    mask = 0 < I[0]
    for i in range(kopt - 1):
        mask *= I[i] < I[i + 1]

    lv = mat.view(-1)

    orders = cand_orders[kopt]

    o = np.array(orders[np.random.randint(len(orders))])
    then = T.zeros((L,) * kopt)
    now = T.zeros_like(then)
    for i in range(kopt):
        now += lv[L * I[i] + I[i]]
        then += lv[L * I[o[i] - 1] + I[o[i + 1] - 2]]

    A = then - now

    A[~mask] = -1001

    topv, topi = A.view(-1).topk(min(A.numel(), topk))
    indices = np.random.randint(topi.shape[0], size=(bs,))
    topv = topv[indices]
    topi = topi[indices]

    orders = [o] * bs

    imod = [(topi // L ** (kopt - 1 - i)) % L for i in range(kopt)]

    return T.stack(imod, -1), topv, orders

def ibis(
    model,
    before,
    sentence,
    after,
    bs,
    topk,
    its,
    patience,
    warminit=False,
    gluemask=None,
    original_dy_dx=None,
    tokenizer=None,
    score_fn="nll",
    batchsize=1,
    original_sentence=None,
):
    device = model.device
    sent = sentence

    padded = T.cat([before, sent, after], 0).unsqueeze(0).to(device)
    padded_s = tokenizer.batch_decode(padded)

    gradloss, sc, wsc, spr = score(
        model, padded, padded_s, original_dy_dx, tokenizer, batchsize=1
    )  # Reorder single sentence!

    if score_fn == "grad":
        orscore = gradloss
    else:
        orscore = sc[0]
    yield orscore

    bestscore = orscore

    bestsc = spr[0]

    lfix, rfix, blanks = before.shape[0] - 1, after.shape[0] - 1, 0

    permsents = [
        T.cat([before, T.from_numpy((sent.numpy())), after], 0) for _ in range(bs)
    ]

    bestmask = np.full(permsents[0].shape, True)

    if gluemask is not None:
        bestmask[lfix + 1 : -rfix - 1] = gluemask

    permmasks = [bestmask.copy() for _ in range(bs)]

    if not warminit:
        seg = list(np.nonzero(bestmask[lfix + 1 : -rfix - 1])[0]) + [len(sent)]
        for b in range(bs):
            perm = np.random.permutation(len(seg) - 1)
            ns = []
            nm = []
            for i in range(len(seg) - 1):
                ns.append(sent[seg[perm[i]] : seg[perm[i] + 1]])
                nm.append(
                    bestmask[lfix + 1 : -rfix - 1][seg[perm[i]] : seg[perm[i] + 1]]
                )
            permsents[b][lfix + 1 : -rfix - 1] = T.cat(ns, 0)
            permmasks[b][lfix + 1 : -rfix - 1] = np.concatenate(nm, 0)

    padded = T.stack(permsents, 0).to(device)
    padded_s = tokenizer.batch_decode(padded)

    bestsent = np.zeros(padded[0].shape)

    bestscore = 1e20

    movetype = "init"
    nch = 0
    candidates = np.array([1] * bs)
    last_imp = 0
    for it in range(IBIS_MAX_STEP):

        gc.collect()

        if it - last_imp > patience:
            break

        gradloss, sc, wsc, spr = score(
            model,
            padded,
            padded_s,
            original_dy_dx,
            tokenizer,
            batchsize=batchsize,
        )
        if score_fn == "grad":
            sc = gradloss
        else:
            sc = sc.numpy()

        if it == 0:
            bestwsc = wsc[0]

        if sc.min() < bestscore:
            if it == 0 or np.any(permsents[sc.argmin()] != bestsent):

                nch += 1

                bestsent, bestscore, bgradloss, bestsc, bestwsc, bestmask = (
                    permsents[sc.argmin()],
                    sc.min(),
                    gradloss[sc.argmin()],
                    spr[sc.argmin()],
                    wsc[sc.argmin()],
                    permmasks[sc.argmin()],
                )
                recon_metrics = eval(
                    ori_single,
                    padded_s[sc.argmin()].split("<|endoftext|>")[1:-1],
                    model,
                    tokenizer,
                )
                wandb.log(
                    {
                        "Coarse: GradLoss": sc.min(),
                        "Coarse: BLEU": recon_metrics["bleu"],
                        "Coarse: NER": recon_metrics["ner"],
                        "Coarse: PPL": recon_metrics["ppl"],
                        "Coarse: GradNorm": sc.min() - ALPHA * recon_metrics["ppl"],
                    },
                    step=it + 1,
                )
                text_table.add_data(
                    "Coarse",
                    it + 1,
                    padded_s[sc.argmin()].split("<|endoftext|>")[1:-1],
                    sc.min(),
                    recon_metrics["ner"],
                    recon_metrics["ppl"],
                    recon_metrics["bleu"],
                    sc.min() - ALPHA * recon_metrics["ppl"],
                )

                print(
                    f"[Progress] Iter: {it}    Score: {sc.min()}   Sentence: {padded_s[sc.argmin()]}  Grad_norm: {sc.min() - ALPHA * recon_metrics['ppl']} {recon_metrics}",
                )

                if type(bestsent) == T.Tensor:
                    bestsent = bestsent.numpy()

                last_imp = it

                yield (
                    it,
                    movetype,
                    bestscore,
                    bgradloss,
                    bestsent,
                    bestmask,
                    recon_metrics,
                )

            if bestscore == 0.0:
                break
        else:
            recon_metrics = eval(
                ori_single,
                padded_s[sc.argmin()].split("<|endoftext|>")[1:-1],
                model,
                tokenizer,
            )
            print(
                f"[Stay] Iter: {it}    Score: {bestscore} Step best: {padded_s[sc.argmin()]}  Grad_norm: {sc.min() - ALPHA * recon_metrics['ppl']}  {recon_metrics}"
            )
            wandb.log(
                {
                    "Coarse: GradLoss": bestscore,
                    "Coarse: BLEU": recon_metrics["bleu"],
                    "Coarse: NER": recon_metrics["ner"],
                    "Coarse: PPL": recon_metrics["ppl"],
                    "Coarse: GradNorm": sc.min() - ALPHA * recon_metrics["ppl"],
                },
                step=it + 1,
            )
            text_table.add_data(
                "Coarse (stay)",
                it + 1,
                padded_s[sc.argmin()].split("<|endoftext|>")[1:-1],
                sc.min(),
                recon_metrics["ner"],
                recon_metrics["ppl"],
                recon_metrics["bleu"],
                sc.min() - ALPHA * recon_metrics["ppl"],
            )

        thespr = bestsc

        kopt = np.random.randint(3, 6)

        cutprobs = np.ones_like(bestwsc)

        cutprobs[~bestmask] = 0.0

        cutprobs[lfix] = 100
        cutprobs[-1 - rfix] = 100

        if it % 2 == 0 and len(bestsent) - lfix - rfix > 6:
            ncand = bestmask[lfix : len(bestsent) - rfix].sum()
            ncand = min(40, ncand)

            l, r = lfix, len(bestsent) - rfix
            candidates = np.random.choice(
                np.arange(l, r),
                replace=False,
                p=cutprobs[l:r] / cutprobs[l:r].sum(),
                size=(ncand,),
            )
            candidates.sort()

            movetype = f"GS {kopt}"

        else:

            ropt = np.random.randint(7, 15)

            try:
                start = np.random.randint(lfix + 1, len(bestsent) - ropt - rfix)

                l, r = start, start + ropt

                candidates = np.random.choice(
                    np.arange(l, r),
                    replace=False,
                    p=cutprobs[l:r] / cutprobs[l:r].sum(),
                    size=(min(ropt, (cutprobs[l:r] > 0).sum()),),
                )

            except:
                ropt = min(15, len(bestsent) - lfix - rfix - 2)
                start = np.random.randint(
                    lfix + 1, max(lfix + 2, len(bestsent) - ropt - rfix)
                )

                l, r = start, start + ropt
                candidates = np.random.choice(
                    np.arange(l, r),
                    replace=False,
                    p=cutprobs[l:r] / cutprobs[l:r].sum(),
                    size=(min(ropt, (cutprobs[l:r] > 0).sum()),),
                )

            candidates.sort()

            movetype = f"LS {kopt}"

        links = thespr[:, bestsent[candidates]][candidates]

        permsents = []
        permmasks = []

        i, v, o = shuffle_proposals(links, topk, bs, kopt)

        for j in range(bs):
            inds = [candidates[0]] + list(candidates[i[j]]) + [candidates[-1]]
            if v[j] > -1000:
                pieces = [bestsent[: inds[0]]]
                maskpieces = [bestmask[: inds[0]]]
                for k in range(kopt + 1):
                    pieces.append(bestsent[inds[o[j][k] - 1] : inds[o[j][k]]])
                    maskpieces.append(bestmask[inds[o[j][k] - 1] : inds[o[j][k]]])
                pieces.append(bestsent[inds[-1] :])
                newsent = np.concatenate(pieces, 0)

                maskpieces.append(bestmask[inds[-1] :])
                newmask = np.concatenate(maskpieces, 0)
            else:
                newsent, newmask = bestsent, bestmask

            permsents.append(newsent)
            permmasks.append(newmask)

        padded = T.stack(list(map(T.from_numpy, permsents)), 0).to(device)
        padded_s = tokenizer.batch_decode(padded)

def createModels():
    import GPUtil
    from transformers import AutoModelForCausalLM, GPT2Tokenizer

    tokenizer_GPT = GPT2Tokenizer.from_pretrained("gpt2")
    tokenizer_GPT.pad_token = tokenizer_GPT.eos_token
    gpt2 = AutoModelForCausalLM.from_pretrained(
        "experiments/wikitext-103/wikitext-103-203456-samples/pretrained_models/batchsize_64/pretrain_epoch_31_iter_98000_linear_1e-05-randinit:False.model",
    )

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

    gpt2 = gpt2.to(device)
    gpt2.eval()

    return tokenizer_GPT, gpt2

def ibisAlgorithm(
    original_dy_dx, initial_sentence, model, tokenizer, attack_bs, original_sentence
):
    # no idea what all these args do
    b = 64
    B = 256
    patience = 128

    vocab = tokenizer.get_vocab()
    vocab = {vocab[i]: i for i in vocab}
    V = len(vocab)
    unbreakable = np.zeros((V,))
    for v in range(V):
        unbreakable[v] = vocab[v][0].lower() in "abcdefghijklmnopqrstuvwxyz"

    sentence = torch.LongTensor(tokenizer.encode(initial_sentence))
    before = torch.LongTensor(
        tokenizer.encode("<|endoftext|>")
    )  # ],return_tensors='pt').input_ids[0]
    after = torch.LongTensor(
        tokenizer.encode("<|endoftext|>")
    )  # ],return_tensors='pt').input_ids[0]

    mask = 1 - unbreakable[sentence]
    mask[0] = 1

    df = pd.DataFrame(columns=["round", "mode", "nll", "gradloss", "sentence"])
    for nch, k in enumerate(
        ibis(
            model,
            before,
            sentence,
            after,
            b,
            B,
            IBIS_MAX_STEP,
            patience,
            warminit=True,
            gluemask=mask,
            original_dy_dx=original_dy_dx,
            tokenizer=tokenizer,
            score_fn="grad",
            batchsize=attack_bs,
            original_sentence=original_sentence,
        )
    ):
        # if warminit is set to False, the ibis will start with random shuffles of the initialized sentence (that is to say, if we already have a good init, we should set warminit=True)
        if nch == 0:
            starting = k.item()
        else:
            df = df.append(
                {
                    "round": k[0],
                    "mode": k[1],
                    "nll": k[2],
                    "gradloss": k[3],
                    "sentence": tokenizer.decode(k[4][1:-1]),
                },
                ignore_index=True,
            )
            ibis_recon_metrics = k[6]
    return df[df["gradloss"] == df["gradloss"].min()], ibis_recon_metrics

def do_leakage(tokenizer, gpt2, original_sentence, ori_single, bs_recon, seed=1):
    import ast
    import re
    from functools import reduce

    np.random.seed(seed)
    torch.manual_seed(seed)

    bs_dy_dx = get_grad_gpt2(bs_recon, tokenizer, gpt2)
    ori_dy_dx = [get_grad_gpt2([o], tokenizer, gpt2) for o in original_sentence[:3]]
    ori_logs = [eval(ori_single, [o], gpt2, tokenizer) for o in original_sentence]
    bs_recon_metrics = eval(ori_single, bs_recon, gpt2, tokenizer)
    bs_strip = [
        "".join(re.split(r'( *[\.\?!][\'"\)\]]* *)', bs_recon[0])[:2])
    ]  # strip the first generated part

    bs_strip_dy_dx = get_grad_gpt2(bs_strip, tokenizer, gpt2)
    bs_strip_metrics = eval(ori_single, bs_strip, gpt2, tokenizer)

    # print(
    #     original_sentence,
    #     "GradNorm:",
    #     [torch.norm(odydx[LAYER_IDX]).cpu().numpy().item() for odydx in ori_dy_dx],
    #     "PPL:",
    #     [ologs["ppl"] for ologs in ori_logs],
    # )
    # print(
    #     bs_recon,
    #     "GradNorm:",
    #     torch.norm(bs_dy_dx[LAYER_IDX]).cpu().numpy().item(),
    #     "PPL:",
    #     bs_recon_metrics["ppl"],
    #     "BLEU:",
    #     bs_recon_metrics["bleu"],
    # )
    # print(
    #     bs_strip,
    #     "GradNorm:",
    #     torch.norm(bs_strip_dy_dx[LAYER_IDX]).cpu().numpy().item(),
    #     "PPL:",
    #     bs_strip_metrics["ppl"],
    #     "BLEU:",
    #     bs_strip_metrics["bleu"],
    # )

    if bs_strip_metrics["ppl"] < bs_recon_metrics["ppl"]:
        best_beamsearch = bs_strip
    else:
        best_beamsearch = bs_recon

    beamsearch_dy_dx = get_grad_gpt2(best_beamsearch, tokenizer, gpt2)
    if len(original_sentence) <= 64:
        original_dy_dx = get_grad_gpt2(original_sentence, tokenizer, gpt2)
        beamsearch_gradloss = (
            calculateGradLoss(original_dy_dx, beamsearch_dy_dx)["MSE"]
            .detach()
            .cpu()
            .numpy()
        ).item()
    else:
        original_dy_dx = None
        beamsearch_gradloss = 1e10

    wandb.log(
        {
            "Coarse: GradLoss": beamsearch_gradloss,
            "Coarse: BLEU": bs_recon_metrics["bleu"],
            "Coarse: NER": bs_recon_metrics["ner"],
            "Coarse: PPL": bs_recon_metrics["ppl"],
        },
        step=0,
    )
    text_table.add_data("Ori", 0, original_sentence, None, None, None, None, None)
    text_table.add_data(
        "Coarse",
        0,
        best_beamsearch,
        beamsearch_gradloss,
        bs_recon_metrics["ner"],
        bs_recon_metrics["ppl"],
        bs_recon_metrics["bleu"],
        None,
    )

    original_sentence_tokens = reduce(
        lambda x, y: x + y, [tokenizer.encode(s) for s in original_sentence]
    )
    best_beamsearch_tokens = reduce(
        lambda x, y: x + tokenizer.encode("<|endoftext|>") + y,
        [tokenizer.encode(s) for s in best_beamsearch],  # seperate sentences
    )

    best_ibis = best_beamsearch
    ibis_recon_metrics = bs_recon_metrics
    ibis_gradloss = beamsearch_gradloss

    best_local = best_beamsearch
    local_recon_metrics = bs_recon_metrics
    local_gradloss = beamsearch_gradloss

    #### IBIS ####
    if IBIS_MAX_STEP > 0:
        best_beamsearch_cat = tokenizer.decode(best_beamsearch_tokens)
        ibis_result, ibis_recon_metrics = ibisAlgorithm(
            original_dy_dx,
            best_beamsearch_cat,
            gpt2,
            tokenizer,
            attack_bs=len(best_beamsearch),
            original_sentence=original_sentence,
        )
        best_ibis = ibis_result["sentence"].item()
        best_ibis_sentences = best_ibis.split("<|endoftext|>")
        ibis_dy_dx = get_grad_gpt2(best_ibis_sentences, tokenizer, gpt2)
        if original_dy_dx is not None:
            ibis_gradloss = (
                calculateGradLoss(original_dy_dx, ibis_dy_dx)["MSE"]
                .detach()
                .cpu()
                .numpy()
            ).item()
        else:
            ibis_gradloss = 1e10

        # terminate early if the ibis algorithm perfectly recovers the original sentence
        if best_ibis == original_sentence:
            global LOCAL_MAX_STEP
            LOCAL_MAX_STEP = (
                0  # Do not run local swapping if global swapping gets perfect results
            )

        tokenized_ibis = tokenizer(best_ibis)["input_ids"]
        losses = [ibis_gradloss]
    else:
        tokenized_ibis = best_beamsearch_tokens
        losses = [beamsearch_gradloss]

    best_local = best_ibis
    local_recon_metrics = ibis_recon_metrics
    local_gradloss = ibis_gradloss

    ### LOCAL ###
    localswap = LocalSwap(gpt2, tokenizer)
    sep_idx = np.where(np.asarray(tokenized_ibis) == 50256)[0].tolist()
    # print(sep_idx)
    reorderGenerator = localswap.createAllOrderings(
        len(tokenized_ibis), without_order=True, sep_idx=sep_idx
    )

    current_tokens = tokenized_ibis
    all_tokens = original_sentence_tokens
    if args.insert:
        insertGenerator = localswap.createAllOrderings(len(all_tokens), sep_idx=sep_idx)
    else:
        insertGenerator = None

    last_imp = 0
    for i in range(1, LOCAL_MAX_STEP):
        reordered = localswap.reorderAlgorithm(
            original_dy_dx, current_tokens, reorderGenerator
        )
        inserted = localswap.insertAlgorithm(
            original_dy_dx, current_tokens, all_tokens, insertGenerator
        )

        best = min(
            [
                reordered,
                inserted,
                localswap.computeStats(current_tokens, original_dy_dx, do_split=True),
            ],
            key=lambda x: x["gradloss_mse"],
        )
        current_tokens = best["tokens"]
        if best["gradloss_mse"] != losses[-1]:
            reorderGenerator = localswap.createAllOrderings(
                len(current_tokens), without_order=True, sep_idx=sep_idx
            )
            if args.insert:
                insertGenerator = localswap.createAllOrderings(
                    len(all_tokens), sep_idx=sep_idx
                )
            else:
                insertGenerator = None

            sys.stdout.flush()
            losses.append(best["gradloss_mse"])

            local_recon_metrics = eval(ori_single, best["sentence"], gpt2, tokenizer)
            print(
                i,
                tokenizer.decode(current_tokens),
                "GradNorm:",
                best["gradloss_mse"] - ALPHA * local_recon_metrics["ppl"],
                "PPL:",
                local_recon_metrics["ppl"],
                "BLEU:",
                local_recon_metrics["bleu"],
            )
            wandb.log(
                {
                    "Fine: GradLoss": best["gradloss_mse"],
                    "Fine: BLEU": local_recon_metrics["bleu"],
                    "Fine: NER": local_recon_metrics["ner"],
                    "Fine: PPL": local_recon_metrics["ppl"],
                    "Fine: GradNorm": best["gradloss_mse"]
                    - ALPHA * local_recon_metrics["ppl"],
                },
                step=i + IBIS_MAX_STEP,
            )
            text_table.add_data(
                "Fine",
                i,
                best["sentence"],
                best["gradloss_mse"],
                local_recon_metrics["ner"],
                local_recon_metrics["ppl"],
                local_recon_metrics["bleu"],
                best["gradloss_mse"] - ALPHA * local_recon_metrics["ppl"],
            )
            last_imp = i
        else:
            wandb.log(
                {
                    "Fine: GradLoss": best["gradloss_mse"],
                    "Fine: BLEU": local_recon_metrics["bleu"],
                    "Fine: NER": local_recon_metrics["ner"],
                    "Fine: PPL": local_recon_metrics["ppl"],
                    "Fine: GradNorm": best["gradloss_mse"]
                    - ALPHA * local_recon_metrics["ppl"],
                },
                step=i + IBIS_MAX_STEP,
            )
            text_table.add_data(
                "Fine",
                i,
                best["sentence"],
                best["gradloss_mse"],
                local_recon_metrics["ner"],
                local_recon_metrics["ppl"],
                local_recon_metrics["bleu"],
                best["gradloss_mse"] - ALPHA * local_recon_metrics["ppl"],
            )

        if best["gradloss_mse"] == 0.0:
            print("Finished early due to perfect recovery!")
            break
        if i - last_imp > LOCAL_PATIENCE:
            break

    best_local = best["sentence"]
    local_gradloss = best["gradloss_mse"]

    res = {
        "Original": original_sentence,
        "After_BS": {
            "sent": bs_recon,
            "score": bs_recon_metrics,
            # "gradloss": beamsearch_gradloss,
        },
        "After_strip": {
            "sent": bs_strip,
            "score": bs_strip_metrics,
            # "gradloss": bs_strip_dy_dx,
        },
        "After_coarse": {
            "sent": best_ibis,
            "score": ibis_recon_metrics,
            # "gradloss": ibis_gradloss,
        },
        "After_fine": {
            "sent": best_local,
            "score": local_recon_metrics,
            # "gradloss": local_gradloss,
        },
    }
    return res

def parse_config(row, wandb):
    config = {}
    wandb.config.batch_size = row["batchsize"]
    wandb.config.batch_id = row["sample_num"]
    wandb.config.beam_size = row["sample_num"]
    wandb.config.iter = row["train_iteration"]
    wandb.config.scheduler = row["train_lrsched"]
    wandb.config.randinit = row["randinit"]
    wandb.config.freqknown = row["freqknown"]
    wandb.config.lr = row["train_lr"]
    return config

if __name__ == "__main__":
    IBIS_MAX_STEP = 100
    LOCAL_MAX_STEP = 100
    LOCAL_PATIENCE = 50

    LAYER_IDX = 68  # 0: wte; 1: wpe; 68: h.5.ln_2.weight

    parser = argparse.ArgumentParser()
    parser.add_argument("--bs", type=int, default=4)
    parser.add_argument("--data", type=str, default="wikitext")
    parser.add_argument("--alpha", type=float, default=10)
    parser.add_argument("--freqknown", action="store_true")
    parser.add_argument("--insert", action="store_true")
    parser.add_argument("--delete", action="store_true")
    parser.add_argument("--n_beamheight", action="store_true")
    args = parser.parse_args()

    if args.data == "wikitext":
        if args.n_beamheight:
            args.fname = f"experiments/wikitext-103-lrsched/wikitext-103-203456-samples/bs_results_height_n_freqknown_{args.freqknown}.pkl"
        else:
            args.fname = f"experiments/wikitext-103-lrsched/wikitext-103-203456-samples/bs_results_height_1_freqknown_{args.freqknown}.pkl"
    else:
        if args.n_beamheight:
            args.fname = f"experiments/enron/enron-30000/bs_results_height_n_freqknown_{args.freqknown}.pkl"
        else:
            args.fname = f"experiments/enron/enron-30000/bs_results_height_1_freqknown_{args.freqknown}.pkl"
    ALPHA = args.alpha
    exp = pd.read_pickle(args.fname)

    run_name = f"bs={args.bs}"
    run = wandb.init(
        project="reorder-single",
        name=f"single-{run_name}-ibis-{IBIS_MAX_STEP}-local-{LOCAL_MAX_STEP}-layer-{LAYER_IDX}-insert{args.insert}-delete-{args.delete}",
    )
    text_table = wandb.Table(
        columns=[
            "Mode",
            "Iter",
            "Sentence",
            "GradLoss",
            "NER",
            "PPL",
            "BLEU",
            "GradNorm",
        ]
    )

    rows = exp[exp["batchsize"] == args.bs]

    parse_config(rows, wandb)

    tokenizer, gpt2 = createModels()
    if args.data == "wikitext":
        if args.n_beamheight:
            save_file = f"experiments/wikitext-103-lrsched/wikitext-103-203456-samples/reorder_single_beamheight_n/bs-{args.bs}/freqknown-{args.freqknown}-span-{IBIS_MAX_STEP}-token_{LOCAL_MAX_STEP}-alpha={args.alpha}_insert={args.insert}_delete={args.delete}.json"
        else:
            save_file = f"experiments/wikitext-103-lrsched/wikitext-103-203456-samples/reorder_single_beamheight_1/bs-{args.bs}/freqknown-{args.freqknown}-span-{IBIS_MAX_STEP}-token_{LOCAL_MAX_STEP}-alpha={args.alpha}_insert={args.insert}_delete={args.delete}.json"
    else:
        if args.n_beamheight:
            save_file = f"experiments/enron/enron-30000/reorder_single_beamheight_n/bs-{args.bs}/freqknown-{args.freqknown}-span-{IBIS_MAX_STEP}-token_{LOCAL_MAX_STEP}-alpha={args.alpha}_insert={args.insert}_delete={args.delete}.json"
        else:
            save_file = f"experiments/enron/enron-30000/reorder_single_beamheight_1/bs-{args.bs}/freqknown-{args.freqknown}-span-{IBIS_MAX_STEP}-token_{LOCAL_MAX_STEP}-alpha={args.alpha}_insert={args.insert}_delete={args.delete}.json"
    rows.to_pickle(
        f"experiments/wikitext-103-lrsched/wikitext-103-203456-samples/reorder_single_beamheight_1/bs-{args.bs}/bs_results_height_1_freqknown_False.pkl"
    )  # save orginal sentences

    res = {}
    for ir in range(len(rows)):
        ori = rows["original"].iloc[ir]
        ori_single = [rows["original single"].iloc[ir]]
        bs_single = rows["candidate single"].iloc[ir]

        if rows["bleu"].iloc[ir] == 1:
            res[ir] = {
                "Original": ori,
                "After_BS": {
                    "sent": ori,
                    "score": None,
                    # "gradloss": beamsearch_gradloss,
                },
                "After_strip": {
                    "sent": ori,
                    "score": None,
                    # "gradloss": bs_strip_dy_dx,
                },
                "After_coarse": {
                    "sent": ori,
                    "score": None,
                    # "gradloss": ibis_gradloss,
                },
                "After_fine": {
                    "sent": ori,
                    "score": None,
                    # "gradloss": local_gradloss,
                },
            }
        else:
            res[ir] = do_leakage(tokenizer, gpt2, ori, ori_single, [bs_single], 1)
        run.log({"reordering_samples": text_table})
        # save after each run
        with open(save_file, "w") as f:
            json.dump(res, f, indent=4, separators=(",", ": "))

repo_name: Princeton-SysML/FILM

repo_language: Python


text:
import os
import time
import pathlib
import subprocess
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--datapath", type=pathlib.Path, required=True)

parser.add_argument(
    "--outputpath", type=pathlib.Path, default=pathlib.Path("output.json")
)
parser.add_argument("--modelweights", type=pathlib.Path)

parser.add_argument("--num_beams", type=int, default=32)
parser.add_argument("--ngram_size_penalty", type=int, default=2)
parser.add_argument("--n_repetitions", type=int, default=1)
parser.add_argument("--frequency_known", action="store_true")

args = parser.parse_args()

beamsearch_rust_dir = pathlib.Path(__file__).parent / "beamsearch_rust"
cargo_path = subprocess.run(["which", "cargo"], stdout=subprocess.PIPE).stdout.decode(
    "utf-8"
)[:-1]
process_env = os.environ.copy()

print("Running beamsearch with the following args", args)
process = subprocess.Popen(
    [
        f"{cargo_path} run "
        "--release "
        "--manifest-path "
        f"{(beamsearch_rust_dir / 'Cargo.toml').absolute().as_posix()} "
        "-- "
        "--num-beams "
        f"{args.num_beams} "
        "--datapath "
        f"{args.datapath.absolute().as_posix()} "
        "--outputpath "
        f"{args.outputpath} "
        "--n-repetitions "
        f"{args.n_repetitions} "
        f"{'--no-frequencies' if not args.frequency_known else ''} "
        f"--no-repeat-ngram-size "
        f"{args.ngram_size_penalty} "
        f"{('--modelweights ' + str(args.modelweights)) if args.modelweights else ''} "
    ],
    shell=True,
    stdout=pathlib.Path("stdout.log").open("w"),
    stderr=pathlib.Path("stderr.log").open("w"),
    env=process_env,
)

while True:
    time.sleep(1)
    if process.poll() is not None:
        break

mime_type: text/plain


text:
from collections import Counter
from pprint import pprint

import nltk
import numpy as np
import torch
from nltk.corpus import stopwords
from torch.autograd import grad
from transformers import GPT2LMHeadModel

def print_best(metric, samples, name1, scores1, name2=None, scores2=None, n=10):
    """print the `n` best samples according to the given `metric`"""
    idxs = np.argsort(metric)[::-1][:n]

    for i, idx in enumerate(idxs):
        if scores2 is not None:
            print(
                f"{i+1}: {name1}={scores1[idx]:.3f}, {name2}={scores2[idx]:.3f}, score={metric[idx]:.3f}"
            )
        else:
            print(f"{i+1}: {name1}={scores1[idx]:.3f}, , score={metric[idx]:.3f}")
        pprint(samples[i])

def calculateGradLoss(original_dy_dx, generated_s_dy_dx, loss_fns=["MSE", "cosine"]):
    grad_diff = {"MSE": 0, "cosine": 0}

    for gx, gy in zip(generated_s_dy_dx, original_dy_dx):
        if "MSE" in loss_fns:
            grad_diff["MSE"] += ((gx - gy) ** 2).sum()
        if "cosine" in loss_fns:
            gx = gx.reshape((-1))
            gy = gy.reshape((-1))
            grad_diff["cosine"] += 1 - torch.sum(gx * gy, dim=-1) / (
                torch.norm(gx, dim=-1) * torch.norm(gy, dim=-1) + 0.000001
            )

    return grad_diff

def calculatePerplexity(sentence, model, tokenizer):
    """exp(loss)"""
    input_ids = torch.tensor(tokenizer.encode(sentence)).unsqueeze(0)
    input_ids = input_ids.to(model.device)

    with torch.no_grad():
        outputs = model(input_ids, labels=input_ids)
    loss, logits = outputs[:2]
    ppl1 = torch.exp(loss)

    return ppl1

def get_grad_gpt2(sentence, tokenizer_GPT, gpt2):
    tokens = tokenizer_GPT(
        sentence,
        padding="max_length",
        truncation=True,
        max_length=gpt2.config.max_length,
    )[
        "input_ids"
    ]  # TODO: change max_length

    if not torch.is_tensor(tokens):
        tokens = torch.tensor(tokens).long()
    tokens = tokens.to(gpt2.device)
    input_shape = tokens.size()
    position_ids = torch.arange(
        0, input_shape[-1], dtype=torch.long, device=gpt2.device
    )
    position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])

    gpt2.zero_grad()
    gpt2_outputs = gpt2(
        input_ids=tokens, output_hidden_states=True, return_dict=True, labels=tokens,
    )
    new_dy_dx = grad(
        gpt2_outputs.loss, gpt2.transformer.parameters(), retain_graph=True
    )

    list_dy_dx = list((_.detach().clone() for _ in new_dy_dx))

    return list_dy_dx

def get_grad(bert, sentence, tokenizer, device, pass_token=False, batched=False):
    if pass_token:
        original = sentence
    else:
        if batched:
            original = tokenizer.batch_encode_plus(
                sentence, return_tensors="pt", padding=True, truncation=True
            )["input_ids"].to(device)
        else:
            if isinstance(sentence, str):
                original = tokenizer.encode(sentence, return_tensors="pt").to(device)
            else:
                original = tokenizer.encode(sentence[0], return_tensors="pt").to(device)
    input_shape = original.size()

    attention_mask = torch.ones(input_shape, device=device)
    head_mask = None
    head_mask = bert.get_head_mask(head_mask, bert.config.num_hidden_layers)

    output_attentions = bert.config.output_attentions
    output_hidden_states = bert.config.output_hidden_states
    return_dict = bert.config.use_return_dict

    bert.transformer.zero_grad()

    dlbrt_output = bert.transformer(
        x=bert.embeddings(original),
        attn_mask=attention_mask,
        head_mask=head_mask,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    hidden_states = dlbrt_output[0]

    dy_dx = grad(hidden_states.sum(), bert.transformer.parameters())
    list_dy_dx = list((_.detach().clone() for _ in dy_dx))

    return list_dy_dx

def generate_sentences(
    LMhead: GPT2LMHeadModel,
    tokenizer,
    leak_results_words,
    rep=1,
    max_length=30,
    batch_size=1,
    max_rep=2,
    max_rep_dict=None,
    prompts_list=None,
    leak_results_tokens=None,
    device="cpu",
    temperature=1,
    mode="sample",
    top_p=0.95,
):
    sentences = []
    generated_token_lists = []
    num_parallel = 10

    stop_words = list(stopwords.words("english"))
    stop_words = " ".join(stop_words)
    stop_words_tokens = tokenizer.encode(stop_words, return_tensors="pt").to(device)

    if leak_results_tokens is None:
        leak_results_tokens = tokenizer.encode(
            leak_results_words, return_tensors="pt"
        ).to(device)

    token2word = {}
    max_rep_dict_ori = max_rep_dict
    if max_rep_dict_ori is None:
        max_rep_dict = {}
    for token_id in leak_results_tokens:
        token2word[token_id] = tokenizer.decode(int(token_id))
        if max_rep_dict_ori is None:
            if token_id in stop_words_tokens:
                max_rep_dict[token_id] = (
                    max_rep * 2
                )  ## Yangsibo: we can also customize this value
            else:
                max_rep_dict[token_id] = max_rep
    print(max_rep_dict)
    print(np.sum([v for (k, v) in max_rep_dict.items()]))
    print(token2word)

    def prefix_allowed_tokens_fn(batch_id, sent):
        rep_dict = Counter(sent.cpu().numpy().tolist())
        allowed = Counter(max_rep_dict)
        allowed.subtract(rep_dict)

        ret = [k for (k, v) in allowed.items() if v > 0]

        prompts_ids = [
            tokenizer.encode(prompt, return_tensors="pt").cpu()[0].numpy()[0]
            for prompt in prompts_list
        ]

        if len(ret) == 0:
            return (
                tokenizer.encode("###", return_tensors="pt").to(device).cpu()[0].numpy()
            )  ## TODO: should implement this as a stopping criteria

        # else:
        #     return ret
        # Yangsibo: it seems that this heuristic makes the generation worse
        if (
            len(sent) % max_length == 0
            or sent[-1]
            == tokenizer.encode(".", return_tensors="pt").to(device).cpu()[0].numpy()[0]
        ):  # Only allow prompt words at begining of sequences
            if len(set(prompts_ids) & set(ret)) > 0:
                return list(set(prompts_ids) & set(ret))
            else:
                return ret

        else:
            return list(set(ret) - set(prompts_ids))

    if prompts_list is None:
        prompts_list = leak_results_words  # Use each word in the inferred set as prompt

    for prompts in prompts_list:
        if prompts in ["[CLS]", "[SEP]", " "]:
            continue
        print(f'Using "{prompts}" as prompt ...')
        if mode == "sample":
            for r in range(rep):
                prompts_rep = [prompts for _ in range(num_parallel)]
                inputs = tokenizer(prompts_rep, return_tensors="pt", truncation=True)
                # print(prompts, inputs['input_ids'].shape)

                # Yangsibo: maybe we should tune the hyper-params here
                output_sequences = LMhead.generate(
                    input_ids=inputs["input_ids"].to(device),
                    max_length=max_length * batch_size,
                    min_length=max_length * batch_size,
                    do_sample=True,
                    top_k=max_length,
                    temperature=temperature,
                    top_p=top_p,
                    repetition_penalty=1.2,
                    prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
                    output_scores=True,
                )

                end_sign = tokenizer.encode("###", return_tensors="pt").to(device)[0]
                # import pdb; pdb.set_trace()
                output_sequences = [
                    [token_id for token_id in output_sequence if token_id != end_sign]
                    for output_sequence in output_sequences
                ]
                if batch_size > 1:
                    output_sequences = [
                        shard_sequence(output_sequence, max_length)
                        for output_sequence in output_sequences
                    ]
                    # import pdb; pdb.set_trace()
                    generated_token = [
                        [
                            [tokenizer.decode(int(token_id)) for token_id in s]
                            for s in output_sequence
                        ]
                        for output_sequence in output_sequences
                    ]
                    texts = [
                        tokenizer.batch_decode(
                            output_sequence, skip_special_tokens=False
                        )
                        for output_sequence in output_sequences
                    ]
                else:
                    generated_token = [
                        [
                            tokenizer.decode(int(token_id))
                            for token_id in output_sequence
                        ]
                        for output_sequence in output_sequences
                    ]
                    texts = tokenizer.batch_decode(
                        output_sequences, skip_special_tokens=False
                    )
                print(texts)
                sentences.extend(texts)
                generated_token_lists.extend(generated_token)

        elif mode == "beam":

            num_beams = min(num_parallel * rep, len(leak_results_words))
            num_return = min(num_beams, 50)
            # import pdb; pdb.set_trace()
            inputs = tokenizer(prompts, return_tensors="pt", truncation=True)
            output_sequences = LMhead.generate(
                input_ids=inputs["input_ids"].to(device),
                max_length=max_length * batch_size,
                min_length=max_length * batch_size,
                num_beams=num_beams,
                num_return_sequences=num_return,
                temperature=temperature,
                repetition_penalty=1.2,
                prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            )
            end_sign = tokenizer.encode("###", return_tensors="pt").to(device)[0]
            # import pdb; pdb.set_trace()
            output_sequences = [
                [token_id for token_id in output_sequence if token_id != end_sign]
                for output_sequence in output_sequences
            ]
            if batch_size > 1:
                output_sequences = [
                    shard_sequence(output_sequence, max_length)
                    for output_sequence in output_sequences
                ]
                # import pdb; pdb.set_trace()
                generated_token = [
                    [
                        [tokenizer.decode(int(token_id)) for token_id in s]
                        for s in output_sequence
                    ]
                    for output_sequence in output_sequences
                ]
                texts = [
                    tokenizer.batch_decode(output_sequence, skip_special_tokens=False)
                    for output_sequence in output_sequences
                ]
            else:
                generated_token = [
                    [tokenizer.decode(int(token_id)) for token_id in output_sequence]
                    for output_sequence in output_sequences
                ]
                texts = tokenizer.batch_decode(
                    output_sequences, skip_special_tokens=False
                )
            print(texts)
            sentences.extend(texts)
            generated_token_lists.extend(generated_token)

    return sentences, generated_token_lists

def shard_sequence(input, max_legnth):
    ori_len = len(input)
    ret = []
    for i in range(ori_len // max_legnth):
        s = input[i * max_legnth : (i + 1) * max_legnth]
        ret.append(s)
    return ret

def run_leakage(input, sequences, tokenizer, bert, optim, device):
    input_shape = input.size()
    inputs_embeds = bert.embeddings(input)
    attention_mask = torch.ones(input_shape, device=device)
    head_mask = None
    head_mask = bert.get_head_mask(head_mask, bert.config.num_hidden_layers)

    output_attentions = bert.config.output_attentions
    output_hidden_states = bert.config.output_hidden_states
    return_dict = bert.config.use_return_dict

    optim.zero_grad()
    dlbrt_output = bert.transformer(
        x=inputs_embeds,
        attn_mask=attention_mask,
        head_mask=head_mask,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
    )

    hidden_states = dlbrt_output[0]
    hidden_states.sum().backward()  # can be replaced by any other loss functions

    leaked_seq_length = (
        bert.embeddings.position_embeddings.weight.grad.var(dim=-1).nonzero().max()
    )
    leaked_token_ids = (
        bert.embeddings.word_embeddings.weight.grad.var(dim=-1).nonzero().view(-1)
    )
    leaked_words = tokenizer.decode(leaked_token_ids)  # The order is not preserved.

    s1 = set(leaked_token_ids.tolist())
    s2 = set((k.item() for ids in input for k in ids))
    diff = s1 - s2
    leak_ratio = len(s1) / sum(len(s) for s in input)

    #     print("--" * 40)

    leaked_words = []
    for i in range(leaked_token_ids.size(0)):
        leaked_word = tokenizer.decode(leaked_token_ids[i])
        leaked_words.append(leaked_word)
    #     print("Leaked sentence length:", leaked_seq_length.item() + 1)
    # print("Leaked words:", "|".join(leaked_words))
    #     print("--" * 40)

    original_words = [tokenizer.decode(t) for seq in input for t in seq]
    #     print("Original sentence length:", len(original_words))
    # print("Original sentence:", "|".join(original_words))
    return {
        "leakratio": leak_ratio,
        "originalwords": original_words,
        "leakedwords": leaked_words,
        "leaktokenids": leaked_token_ids,
    }


text:
name: film-nlp
channels:
  - defaults
  - conda-forge
  - pytorch
dependencies:
  - python=3.9
  - mamba
  - pytorch
  - torchvision
  - torchaudio
  - cudatoolkit=11.3
  - numpy
  - pandas
  - seaborn
  - jupyterlab
  - ca-certificates
  - certifi
  - openssl
  - matplotlib
  - nltk
  - tqdm
  - papermill
  - transformers
  - spacy
  - tqdm


text:
import json
import pathlib
from typing import Dict
import numpy as np
import spacy
from rouge_score import rouge_scorer
import re
import pandas as pd
import tqdm
from multiprocessing import Pool, Queue

def namedEntityRatio(nes_ori, nes_cand):
    nes_ori = set(nes_ori)

    if len(nes_ori) == 0:
        return 0
    if len(nes_cand) == 0:
        return 0

    nes_cand = set(nes_cand)

    recall = len(nes_ori.intersection(nes_cand)) / len(nes_ori)
    precision = len(nes_ori.intersection(nes_cand)) / len(nes_cand)
    if recall + precision == 0:
        return 0
    return 2 * (recall * precision) / (recall + precision)

def nerScore(original, generated, nlp):
    cost = np.zeros((len(original), len(generated)))
    named_entities_original = [
        list(map(lambda x: x.text, nlp(o).ents)) for o in original
    ]
    named_entities_candidate = [
        list(map(lambda x: x.text, nlp(o).ents)) for o in generated
    ]
    for i, o in enumerate(named_entities_original):
        for j, g in enumerate(named_entities_candidate):
            cost[i, j] = namedEntityRatio(o, g)

    return cost, named_entities_original, named_entities_candidate

def rougeScores(original, generated, scorer):
    rouge_scores = {}
    for statistic in ["fmeasure", "recall", "precision"]:
        cost_1 = np.zeros((len(original), len(generated)))
        cost_2 = np.zeros((len(original), len(generated)))
        cost_L = np.zeros((len(original), len(generated)))
        for i, o in enumerate(original):
            for j, g in enumerate(generated):
                scores = scorer.score(o, g)
                if statistic == "fmeasure":
                    cost_1[i, j] = scores["rouge1"].fmeasure
                    cost_2[i, j] = scores["rouge2"].fmeasure
                    cost_L[i, j] = scores["rougeL"].fmeasure
                elif statistic == "recall":
                    cost_1[i, j] = scores["rouge1"].recall
                    cost_2[i, j] = scores["rouge2"].recall
                    cost_L[i, j] = scores["rougeL"].recall
                elif statistic == "precision":
                    cost_1[i, j] = scores["rouge1"].precision
                    cost_2[i, j] = scores["rouge2"].precision
                    cost_L[i, j] = scores["rougeL"].precision
        rouge_scores[statistic] = {"rouge1": cost_1, "rouge2": cost_2, "rougeL": cost_L}

    return rouge_scores

def createScorers():
    nlp = spacy.load("en_core_web_sm")
    rouge = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
    return {"nlp": nlp, "rouge": rouge}

def calcMetrics(ref, cand, scorers, strip=True):
    if strip:
        cand = ["".join(re.split(r'( *[\.\?!][\'"\)\]]* *)', c)[:2]) for c in cand]
    ner, ent_orig, ent_cand = nerScore(ref, cand, scorers["nlp"])
    df = pd.DataFrame(
        {
            "original": [ref],
            "candidate": [cand],
            "ner": [ner],
            "ent_orig": [ent_orig],
            "ent_cand": [ent_cand],
        }
    )

    rouge_scores = rougeScores(ref, cand, scorers["rouge"])
    for statistic in ["fmeasure", "recall", "precision"]:
        for rouge in ["1", "2", "L"]:
            df[f"rouge_{rouge}_{statistic}"] = [
                rouge_scores[statistic][f"rouge{rouge}"]
            ]

    return df

def getTopkMetrics(result, params, scorers, k=1):
    topk_data = [
        calcMetrics(result["original"], result["results"][i]["result"], scorers)
        for i in range(0, k)
        if len(result["results"]) > i
    ]
    for idx, metrics in enumerate(topk_data):
        for (k, v) in params.items():
            metrics[k] = v
        metrics["topk"] = idx
    topk_data = pd.concat(topk_data, ignore_index=True)
    return topk_data

def getAllResultsJson(dir: pathlib.Path):
    for child in dir.iterdir():
        if child.name == "result.json":
            params = child.parent / "params.json"
            params = json.load(params.open())
            results = json.load(child.open())
            if params["no_repeat_ngram_size"] is None:
                params["no_repeat_ngram_size"] = 0
            yield {"results": results, "params": params}
        elif child.is_dir():
            yield from getAllResultsJson(child)

def isNewRow(dataframe: pd.DataFrame, parameters: Dict):
    dataframe_columns = set(dataframe.columns)
    for (k, v) in parameters.items():
        if k not in dataframe_columns:
            return True
        dataframe = dataframe[dataframe[k] == v]
    return len(dataframe) == 0

def calculateAllNewMetrics(
    directory: pathlib.Path, existing: pd.DataFrame = None, topk=1
):
    all_results = getAllResultsJson(directory)
    if existing is not None:
        all_results = filter(lambda r: isNewRow(existing, r["params"]), all_results)
    else:
        existing = pd.DataFrame()

    scorers = createScorers()
    dataframes = []
    for result in tqdm.tqdm(list(all_results)):
        dataframes.append(
            getTopkMetrics(result["results"], result["params"], scorers, topk)
        )
    return dataframes

def getBestScores(dataframe: pd.DataFrame, metric: str):
    best_index = [np.unravel_index(np.argmax(v), v.shape) for v in dataframe[metric]]
    best_value = [np.max(v) for v in dataframe[metric]]
    return best_index, best_value

hyperparams = [
    "beamscore-method",
    "num_beams",
    "batchsize",
    "train_iteration",
    "beamgroups",
    "sample_num",
    "beamheight",
    "diversity",
    "randinit",
    "train_lr",
    "train_lrsched",
    "freqknown",
    "no_repeat_ngram_size",
    "dataset_size",
    "n_repetitions",
]

def keepOnlyBestTopk(
    dataframe: pd.DataFrame,
    metric="rouge_1_fmeasure_best_value",
    ascending=False,
    topk=1000,
):
    filtered_hyperparams = [h for h in hyperparams if h in dataframe.columns]
    split_df = dataframe.groupby(filtered_hyperparams)
    rows = []
    for (_, group) in split_df:
        group = group[group["topk"] <= topk]
        rows.append(group.sort_values(metric, ascending=ascending).head(1))
    return pd.concat(rows)

def keepOnlyBestTopkUniqueSentences(
    dataframe: pd.DataFrame,
    metric="rouge_1_fmeasure",
    ascending=False,
    topk=1000,
):
    filtered_hyperparams = [h for h in hyperparams if h in dataframe.columns]
    split_df = dataframe.groupby(filtered_hyperparams)
    rows = []
    for (_, group) in split_df:
        group = group[group["topk"] <= topk]
        split_group = group.groupby(f"{metric}_best_idx")
        for (_, subgroup) in split_group:
            rows.append(subgroup.sort_values(f"{metric}_best_value", ascending=ascending).head(1))
    return pd.concat(rows)

if __name__ == "__main__":
    print(calcMetrics(["The fat cat sits at home."], ["The cat"], createScorers()))

    a, b, c = nerScore(
        [
            "It is variable in form and may be fused dorsally with some of the thoracic segments or occasionally be in two parts, hinged dorsally."
        ],
        [
            "What is known about her is that she had three children: one of whom is unknown;the other two are known not to have been born;and one is not known to be living at all."
        ],
        createScorers()["nlp"],
    )

    print(a, b, c)


text:
# Recovering Private Text in Federated Learning of Language Models

Samyak Gupta*, Yangsibo Huang*, Zexuan Zhong, Tianyu Gao, Kai Li, Danqi Chen (*: equal contribution)
>
> Federated learning allows distributed users to collaboratively train a model while keeping each user's data private. Recently, a growing body of work has demonstrated that an eavesdropping attacker can effectively recover image data from gradients transmitted during federated learning. However, little progress has been made in recovering text data. In this paper, we present a novel attack method FILM for federated learning of language models -- for the first time, we show the feasibility of recovering text from large batch sizes of up to 128 sentences. Different from image-recovery methods which are optimized to match gradients, we take a distinct approach that first identifies a set of words from gradients and then directly reconstructs sentences based on beam search and a prior-based reordering strategy. The key insight of our attack is to leverage either prior knowledge in pre-trained language models or memorization during training. Despite its simplicity, we demonstrate that FILM can work well with several large-scale datasets -- it can extract single sentences with high fidelity even for large batch sizes and recover multiple sentences from the batch successfully if the attack is applied iteratively. We hope our results can motivate future work in developing stronger attacks as well as new defense methods for training language models in federated learning.

## Description

Official implementation of our NeurIPS 2022 paper [Recovering Private Text in Federated Learning of Language Models](https://arxiv.org/pdf/2205.08514.pdf), where we demonstrate a new method (FILM) in which a participant of federated learning (i.e. a curious server, a malicious client, or an observer of transmitted parameters) may leverage memorization in large language models to reconstruct private client data. We evaluate our approach on Wikitex-103 and the Enron Emails datasets.

If you have any questions related to the code or the paper, feel free to email Samyak (samyakg@princeton.edu) and Yangsibo (yangsibo@princeton.edu). If you encounter any problems when using the code, or want to report a bug, you can open an issue. Please try to specify the problem with details so we can help you better and quicker!

## Table of Contents

- [Recovering Private Text in Federated Learning of Language Models](#recovering-private-text-in-federated-learning-of-language-models)
  * [Description](#description)
  * [Setup](#setup)
  * [Datasets](#datasets)
  * [Running Beam Search](#running-beam-search)
  * [Sentence Reordering](#sentence-reordering)
  * [Pre-trained Models](#pre-trained-models)
  * [Results](#results)
  * [Citation](#citation)

Table of contents generated with markdown-toc

## Setup

To install python requirements:

```bash
conda env create -f environment.yml
conda activate film-nlp
python -m spacy download en_core_web_sm
```

To set up beamsearch code:

First [install Rust](https://www.rust-lang.org/tools/install)

Then, run code as follows:

```bash
cd beamsearch_rust
cargo run --release
```

By default, the beamsearch runs only only CPU. To run with CUDA support, set the `TORCH_CUDA_VERSION` environment variable to `cu113` ([See this link for more information](https://github.com/LaurentMazare/tch-rs))

## Datasets

Our experiments used data from the [Enron Emails Dataset](https://www.cs.cmu.edu/~./enron/) and [Wikitext-103](https://huggingface.co/datasets/wikitext). For convenience, we have included a copy of Enron Emails and a subset of Wikitext-103 in the [experiments folder](experiments).

## Running Beam Search

Beam search can be ran in stand-alone form by running the following command:

```setup
cd beamsearch_rust
cargo run --release -- --datapath <DATAPATH> --outputpath <OUTPUTPATH>
```

Data passed to `datapath` must be a json file of a specific form - see [an example](example_data.json) input for a batch of size 16.

By default, beamsearch will be ran with a pretrained GPT2 model from HuggingFace. For a full list of arguments available during beamsearch, you may run the following command:

```bash
cargo run --release -- --help
```

To run with an existing GPT2 model, you must first convert it to a compatible form using [this tool](https://github.com/guillaume-be/rust-bert/blob/master/utils/convert_model.py). Simply save the model (e.g. using model.save_pretrained in HuggingFace), and then run `python convert_model.py pytorch_model.bin` to convert. Then you may use the generated `rust_model.ot` file during beamsearch as follows:

```bash
cargo run --release -- --modelweights <PATH_TO_RUST_MODEL>
```

We also include a [Python frontend](run_beamsearch.py) which passes down args to the rust beamsearch implementation.
## Sentence Reordering

We also include the code used for sentence reordering in [reorder.py](reorder.py). It can be run as follows:

```bash
python reorder_single.py --bs BATCH_SIZE --alpha 1
```

## Pre-trained Models

We currently do not provide any models fine-tuned on the datasets (will be added to the repository at a later data).

For now, you may use pre-trained models from [HuggingFace](https://huggingface.co/models) and fine-tune on the provided datasets.

## Results

Recovery performance for various batch sizes on WikiText-103 (a) and the Enron dataset (b). Solid
lines indicate the average F-Scores of recoveries out of 20 tested mini-batches. The attack is overall weaker for
larger batch sizes (see Table 7 for quantitative results). Moreover, the attack is stronger for the Enron dataset.

## Citation

If you use this code in your research, please cite the following work:
```bibtex
@inproceedings{gupta2022recovering,
   title={Recovering Private Text in Federated Learning of Language Models},
   author={Gupta, Samyak and Huang, Yangsibo and Zhong, Zexuan and Gao, Tianyu and Li, Kai and Chen, Danqi},
   booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
   year={2022}
}
}
```


text:
import logging

logging.basicConfig(level="ERROR")

import argparse
import numpy as np
from pprint import pprint
import sys
import torch
import zlib
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from tqdm import tqdm

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

def calculatePerplexity(sentence, model, tokenizer):
    """
    exp(loss)
    """
    input_ids = torch.tensor(tokenizer.encode(sentence)).unsqueeze(0)
    input_ids = input_ids.to(device)
    with torch.no_grad():
        outputs = model(input_ids, labels=input_ids)
    loss, logits = outputs[:2]
    return torch.exp(loss)

def print_best(metric, samples, name1, scores1, name2=None, scores2=None, n=10):
    """
    print the `n` best samples according to the given `metric`
    """
    idxs = np.argsort(metric)[::-1][:n]

    for i, idx in enumerate(idxs):
        if scores2 is not None:
            print(
                f"{i+1}: {name1}={scores1[idx]:.3f}, {name2}={scores2[idx]:.3f}, score={metric[idx]:.3f}"
            )
        else:
            print(f"{i+1}: {name1}={scores1[idx]:.3f}, , score={metric[idx]:.3f}")

        print()
        # for line in samples[i].split("\n"):
        #    print(f"\t {line.rstrip()}")
        pprint(samples[i])
        print()
        print()

def parse_commoncrawl(wet_file):
    """
    Quick and ugly parsing of a WET file.
    Tested for the May 2021 crawl.
    """
    with open(wet_file) as f:
        lines = f.readlines()

    start_idxs = [i for i in range(len(lines)) if "WARC/1.0" in lines[i]]

    all_eng = ""

    count_eng = 0
    for i in range(len(start_idxs) - 1):
        start = start_idxs[i]
        end = start_idxs[i + 1]
        if "WARC-Identified-Content-Language: eng" in lines[start + 7]:
            count_eng += 1
            for j in range(start + 10, end):
                all_eng += lines[j]

    return all_eng

def main():
    print(f"using device: {device}")

    if args.internet_sampling:
        print("Loading common crawl...")
        cc = parse_commoncrawl(args.wet_file)

    # number of tokens to generate
    seq_len = 256

    # sample from the top_k tokens output by the model
    top_k = 40

    print("Loading GPT2...")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    tokenizer.padding_side = "left"
    tokenizer.pad_token = tokenizer.eos_token

    model1 = GPT2LMHeadModel.from_pretrained("gpt2-xl", return_dict=True).to(device)
    model1.config.pad_token_id = model1.config.eos_token_id
    model2 = GPT2LMHeadModel.from_pretrained("gpt2", return_dict=True).to(device)
    model1.eval()
    model2.eval()

    samples = []
    scores = {"XL": [], "S": [], "Lower": [], "zlib": []}

    num_batches = int(np.ceil(args.N / args.batch_size))
    with tqdm(total=args.N) as pbar:
        for i in range(num_batches):
            # encode the prompts
            if args.internet_sampling:
                # pick a random 10-token prompt in common crawl

                input_len = 10
                input_ids = []
                attention_mask = []

                while len(input_ids) < args.batch_size:
                    # take some random words in common crawl
                    r = np.random.randint(0, len(cc))
                    prompt = " ".join(cc[r : r + 100].split(" ")[1:-1])

                    # make sure we get the same number of tokens for each prompt to enable batching
                    inputs = tokenizer(
                        prompt,
                        return_tensors="pt",
                        max_length=input_len,
                        truncation=True,
                    )
                    if len(inputs["input_ids"][0]) == input_len:
                        input_ids.append(inputs["input_ids"][0])
                        attention_mask.append(inputs["attention_mask"][0])

                inputs = {
                    "input_ids": torch.stack(input_ids),
                    "attention_mask": torch.stack(attention_mask),
                }

                # the actual truncated prompts
                prompts = tokenizer.batch_decode(
                    inputs["input_ids"], skip_special_tokens=True
                )
            else:
                prompts = ["<|endoftext|>"] * args.batch_size
                input_len = 1
                inputs = tokenizer(prompts, return_tensors="pt", padding=True)

            # batch generation
            output_sequences = model1.generate(
                input_ids=inputs["input_ids"].to(device),
                attention_mask=inputs["attention_mask"].to(device),
                max_length=input_len + seq_len,
                do_sample=True,
                top_k=top_k,
                top_p=1.0,
            )

            texts = tokenizer.batch_decode(output_sequences, skip_special_tokens=True)

            for text in texts:
                # perplexity of GPT2-XL and GPT2-S
                p1 = calculatePerplexity(text, model1, tokenizer)
                p2 = calculatePerplexity(text, model2, tokenizer)

                # perplexity on lower-case sample
                p_lower = calculatePerplexity(text.lower(), model1, tokenizer)

                # Zlib "entropy" of sample
                zlib_entropy = len(zlib.compress(bytes(text, "utf-8")))

                samples.append(text)
                scores["XL"].append(p1)
                scores["S"].append(p2)
                scores["Lower"].append(p_lower)
                scores["zlib"].append(zlib_entropy)

            pbar.update(args.batch_size)

    scores["XL"] = np.asarray(scores["XL"])
    scores["S"] = np.asarray(scores["S"])
    scores["Lower"] = np.asarray(scores["Lower"])
    scores["zlib"] = np.asarray(scores["zlib"])

    # Sort by perplexity
    metric = -np.log(scores["XL"])
    print(f"======== top sample by XL perplexity: ========")
    print_best(metric, samples, "PPL", scores["XL"])
    print()
    print()

    # Sort by ratio of log perplexities of S and XL models
    metric = np.log(scores["S"]) / np.log(scores["XL"])
    print(f"======== top sample by ratio of S and XL perplexities: ========")
    print_best(metric, samples, "PPL-XL", scores["XL"], "PPL-S", scores["S"])
    print()
    print()

    # Sort by ratio of log perplexities of lower-case and normal-case perplexities
    metric = np.log(scores["Lower"]) / np.log(scores["XL"])
    print(
        f"======== top sample by ratio of lower-case and normal-case perplexities: ========"
    )
    print_best(metric, samples, "PPL-XL", scores["XL"], "PPL-XL-Lower", scores["Lower"])
    print()
    print()

    # Sort by ratio of Zlib entropy and XL perplexity
    metric = scores["zlib"] / np.log(scores["XL"])
    print(f"======== top sample by ratio of Zlib entropy and XL perplexity: ========")
    print_best(metric, samples, "PPL-XL", scores["XL"], "Zlib", scores["zlib"])

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--N", type=int, default=1000, help="Number of samples to generate"
    )
    parser.add_argument(
        "--batch-size", type=int, default=10, help="Batch size for generation"
    )
    parser.add_argument(
        "--internet-sampling",
        action="store_true",
        help="condition the generation using commoncrawl",
    )
    parser.add_argument(
        "--wet-file", type=str, default=None, help="path to a commoncrawl WET file"
    )
    return parser.parse_args(argv)

if __name__ == "__main__":
    args = parse_arguments(sys.argv[1:])
    main()


text:
[package]
edition = "2021"
name = "beamsearch_rust"
version = "0.1.0"

[profile.dev]
debug = true
split-debuginfo = "packed"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[profile.release]
opt-level = 3

[dependencies]
clap = {version = "3.1.3", features = ["derive"]}
csv = "1.1.6"
either = "1.6.1"
flexbuffers = "2.0.0"
itertools = "0.10.3"
rand = "0.8.5"
rayon = "1.5.1"
rmp-serde = "1.0.0"
rust-bert = "0.17.0"
rust_tokenizers = "7.0.1"
serde = "1.0.136"
serde_json = "1.0.79"
substring = "1.4.5"
tch = "0.6.1"
time = "0.3.7"


text:
use std::{
    cmp::{self, min, Ordering},
    collections::{HashMap, HashSet},
    error::Error,
    ops::{Add, Mul},
};

use crate::util::{get_logits, GPT2ModelAndTokenizer};
use clap::ArgEnum;
use itertools::{izip, Itertools};
use rand::{prelude::SliceRandom, thread_rng};
use rust_tokenizers::tokenizer::TruncationStrategy;
use serde::Serialize;
use tch::{Device, IndexOp, Kind, Tensor};

#[derive(Debug)]
pub(crate) struct Beam {
    pub(crate) state: Box<Tensor>,
    pub(crate) score: f64,
    pub(crate) perplexities: Vec<f64>,
    pub(crate) lengths: Vec<usize>,
    pub(crate) seen_ngrams: Option<HashMap<Vec<i64>, HashSet<i64>>>,
}

impl Clone for Beam {
    fn clone(&self) -> Self {
        Self {
            state: Box::new(self.state.shallow_clone()),
            score: self.score,
            lengths: self.lengths.clone(),
            perplexities: self.perplexities.clone(),
            seen_ngrams: self.seen_ngrams.clone(),
        }
    }
}

impl Beam {
    pub(crate) fn advance(
        &self,
        sent_idx: usize,
        token: i64,
        score: f64,
        perplexity: f64,
        no_repeat_ngram: Option<Vec<i64>>,
    ) -> Beam {
        let mut state = self.state.copy();
        let _ = state
            .i(sent_idx as i64)
            .i(self.lengths[sent_idx] as i64)
            .fill_(token);

        let mut lengths = self.lengths.clone();
        lengths[sent_idx] += 1;

        let mut new_perplexities = self.perplexities.clone();
        new_perplexities[sent_idx] += perplexity;

        let mut new_ngrams = self.seen_ngrams.as_ref().cloned();

        if let Some(n) = no_repeat_ngram {
            new_ngrams
                .as_mut()
                .unwrap()
                .entry(n)
                .or_default()
                .insert(token);
        }

        Beam {
            state: Box::new(state),
            score,
            lengths,
            perplexities: new_perplexities,
            seen_ngrams: new_ngrams,
        }
    }

    pub(crate) fn get_sent_last_n_toks(&self, sent_idx: usize, n_toks: usize) -> Option<Vec<i64>> {
        let length = self.lengths[sent_idx];
        if length >= n_toks {
            let suffix = self
                .state
                .i(sent_idx as i64)
                .slice(0, (length - n_toks) as i64, length as i64, 1)
                .data()
                .into();
            return Some(suffix);
        }
        return None;
    }

    pub(crate) fn get_sent_vec(&self, idx: usize) -> Vec<i64> {
        self.state.i(idx as i64).data().into()
    }

    pub(crate) fn convert_to_string(&self, model: &GPT2ModelAndTokenizer) -> Vec<String> {
        (0..self.lengths.len())
            .map(|idx| self.get_sent_vec(idx))
            .map(|v| model.tokenizer.decode(&v, true, true))
            .collect()
    }

    fn get_remaining_token_counts(&self, token_counts: &HashMap<i64, i64>) -> HashMap<i64, i64> {
        let mut allowed_toks = token_counts.clone();
        for idx in 0..self.lengths.len() {
            let vector = self.get_sent_vec(idx);
            for val in vector {
                *allowed_toks.entry(val).or_insert(0) -= 1;
            }
        }

        allowed_toks
            .iter()
            .filter(|(&_k, &v)| v > 0)
            .map(|(&k, &v)| (k, v))
            .collect()
    }
    // Returns remaining tokens for a beam search state
    fn get_allowed_tokens(&self, token_counts: &HashMap<i64, i64>) -> HashSet<i64> {
        let token_counts = self.get_remaining_token_counts(token_counts);
        let set: HashSet<i64> = HashSet::from_iter(
            token_counts
                .iter()
                .filter(|(_k, &v)| v > 0)
                .map(|(&k, _v)| k),
        );
        set
    }

    fn get_sum_ppl(&self) -> f64 {
        self.perplexities.iter().sum::<f64>()
    }
}

impl PartialOrd for Beam {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        self.score.partial_cmp(&other.score)
    }
}

impl PartialEq for Beam {
    fn eq(&self, other: &Self) -> bool {
        (0..self.lengths.len())
            .map(|idx| self.get_sent_vec(idx) == other.get_sent_vec(idx))
            .all(|x| x)
    }
}

impl Eq for Beam {}

impl Ord for Beam {
    fn cmp(&self, other: &Self) -> Ordering {
        self.score.partial_cmp(&other.score).unwrap()
    }
}

#[derive(Debug, Clone)]
pub(crate) struct BeamGroup {
    beams: Vec<Beam>,
}

/// Computes Hamming diversity between two beam groups.
/// Hamming diversity for two beams A,B is given by minimum number of token substitutions needed to convert from A to B.
/// Each element in the returned tensor is the sum of Hamming diversity for the beam

#[derive(Clone)]
pub(crate) struct BeamSearchOptions {
    pub(crate) max_length: usize,
    pub(crate) num_beams: usize,
    pub(crate) num_groups: usize,
    pub(crate) diversity_penalty: f64,
    pub(crate) max_device_beams: usize,
    pub(crate) frequencies: HashMap<i64, i64>,
    pub(crate) freqknown: bool,
    pub(crate) use_start_tok_freqs: bool,
    pub(crate) beam_height: usize,
    pub(crate) n_repetitions: usize,
    pub(crate) beam_perplexity_scoring: BeamScoreReduction,
    pub(crate) no_repeat_ngram_size: Option<usize>,
}

#[derive(Serialize, Debug)]
struct SingleBeamResult {
    result: Vec<String>,
    perplexity: f64,
    beam_score: f64,
}

#[derive(Serialize, Debug)]
pub(crate) struct BeamSearchResults {
    original: Vec<String>,
    results: Vec<SingleBeamResult>,
}

pub(crate) fn get_frequencies(
    gpt2: &GPT2ModelAndTokenizer,
    sentences: &[String],
    max_length: usize,
) -> std::collections::HashMap<i64, i64> {
    let tokenized_inputs =
        gpt2.tokenizer
            .encode_list(sentences, max_length, &TruncationStrategy::LongestFirst, 1);

    let mut token_counts: HashMap<i64, i64> = HashMap::new();

    for data in tokenized_inputs {
        for id in data.token_ids {
            *token_counts.entry(id).or_insert(0) += 1;
        }
    }
    token_counts
}

#[derive(Debug, Clone, Copy, ArgEnum)]
pub enum BeamScoreReduction {
    Sum,
    Min,
    MinAvg,
}

fn reduce_beam_perplexity(scoring_method: &BeamScoreReduction, beam: &Beam) -> f64 {
    match scoring_method {
        BeamScoreReduction::Sum => beam.perplexities.iter().sum(),
        BeamScoreReduction::Min => beam.perplexities.iter().cloned().reduce(f64::min).unwrap(),
        BeamScoreReduction::MinAvg => izip!(
            beam.lengths.iter().cloned(),
            beam.perplexities.iter().cloned()
        )
        .map(|(l, p)| p / l as f64)
        .reduce(f64::min)
        .unwrap(),
    }
}

pub(crate) fn order_words(
    gpt2: &GPT2ModelAndTokenizer,
    sentences: Vec<String>,
    options: BeamSearchOptions,
) -> Result<BeamSearchResults, Box<dyn Error>> {
    let batch_size = sentences.len();

    let tokenized_inputs = gpt2.tokenizer.encode_list(
        &sentences,
        options.max_length,
        &TruncationStrategy::LongestFirst,
        1,
    );

    let max_token_len = tokenized_inputs
        .iter()
        .map(|v| v.token_ids.len())
        .max()
        .unwrap();
    // let max_token_len = options.max_length;
    println!("Using max sentence length {max_token_len}!");

    let input_ids = Tensor::ones(
        &[sentences.len() as i64, (max_token_len + 1) as i64],
        (Kind::Int64, Device::Cpu),
    )
    .fill_(50256);

    // construct input tensor + count tokens
    for (sent_idx, data) in tokenized_inputs.iter().enumerate() {
        for (token_idx, id) in data.token_ids.iter().enumerate() {
            let _ = input_ids
                .get(sent_idx.try_into()?)
                .get(token_idx.try_into()?)
                .fill_(*id);
        }
    }

    // filter out possible start tokens
    let mut start_tokens: Vec<i64> = vec![];
    for (&k, &count) in &options.frequencies {
        let decoded = gpt2.tokenizer.decode(&[k], true, false);
        if decoded.chars().next().unwrap().is_uppercase() {
            let toks = match options.use_start_tok_freqs {
                true => vec![k].repeat(count as usize),
                false => vec![k],
            };
            start_tokens.extend(toks);
        }
    }

    if start_tokens.len() < batch_size {
        let mult = batch_size as f64 / start_tokens.len() as f64;
        let mult = mult.ceil() as usize;
        println!("Extending starting tokens list with multiplicity {mult}");
        let tmp = start_tokens.clone();
        let iter = tmp.iter().cycle().take(start_tokens.len() * mult);
        start_tokens.extend(iter);
    }

    println!("Found {} possible starting tokens!", start_tokens.len());

    let mut beam_results: Vec<Beam> = Vec::new();
    for _iter in 0..options.n_repetitions {
        let mut beams = beamsearch(gpt2, &start_tokens, &options, &beam_results);
        beam_results.append(&mut beams);
    }

    let mut results: Vec<SingleBeamResult> = beam_results
        .drain(..)
        .map(|b| SingleBeamResult {
            result: b.convert_to_string(gpt2),
            perplexity: reduce_beam_perplexity(&options.beam_perplexity_scoring.clone(), &b),
            beam_score: b.score,
        })
        .collect();
    results.sort_unstable_by(|a, b| b.perplexity.partial_cmp(&a.perplexity).unwrap());

    Ok(BeamSearchResults {
        original: sentences,
        results,
    })
}

fn beamsearch(
    gpt2: &GPT2ModelAndTokenizer,
    initial_beam_states: &[i64],
    options: &BeamSearchOptions,
    previous_results: &Vec<Beam>,
) -> Vec<Beam> {
    tch::Cuda::cudnn_set_benchmark(true);
    let _guard = tch::no_grad_guard();

    let max_sentence_length = min(
        options.frequencies.values().sum(),
        options.max_length as i64,
    );
    let max_iterations: i64 = min(
        options.frequencies.values().sum(),
        (options.max_length * options.beam_height) as i64,
    ) - (options.beam_height as i64);

    let mut init_beam_states = initial_beam_states.to_owned();
    init_beam_states.shuffle(&mut thread_rng());
    let initial_beam_states: Vec<Vec<i64>> = init_beam_states
        .chunks_exact(options.beam_height)
        .map(|x| x.to_vec())
        .collect();

    let initial_beam_tensors: Vec<Tensor> = initial_beam_states
        .iter()
        .map(|v| {
            Tensor::stack(
                &v.iter()
                    .map(|&x| {
                        Tensor::of_slice(
                            &[vec![x], vec![50256; max_sentence_length as usize]].concat(),
                        )
                    })
                    .collect::<Vec<Tensor>>(),
                0,
            )
        })
        .collect();

    let initial_beams: Vec<Beam> = initial_beam_tensors
        .iter()
        .map(|t| Beam {
            state: Box::new(t.shallow_clone()),
            score: 0.0,
            perplexities: vec![0.0; options.beam_height],
            lengths: vec![1; options.beam_height],
            seen_ngrams: options.no_repeat_ngram_size.and(Some(HashMap::new())),
        })
        .collect();

    let mut beam_groups: Vec<BeamGroup> = vec![BeamGroup {
        beams: initial_beams,
    }];

    beam_groups.extend(
        beam_groups
            .iter()
            .cycle()
            .take(options.num_groups - beam_groups.len())
            .cloned()
            .collect::<Vec<BeamGroup>>(),
    );

    let mut beam_s_time = time::Instant::now();

    // Penalize previously seen n-grams
    let previously_seen_ngrams = previous_results
        .iter()
        .map(|b| b.seen_ngrams.as_ref().unwrap())
        .cloned()
        .reduce(|mut a, b| {
            b.iter().for_each(|(k, v)| {
                a.entry(k.to_owned()).or_default().extend(v.iter());
            });
            a
        })
        .unwrap_or_default();

    for iteration in 1..=max_iterations {
        let new_s_time = time::Instant::now();
        let remaining_tokens: i64 = beam_groups[0].beams[0]
            .get_remaining_token_counts(&options.frequencies)
            .values()
            .sum();
        println!(
            "iteration {}/{} {:?} (remaining: {})",
            iteration,
            max_iterations,
            new_s_time - beam_s_time,
            remaining_tokens,
        );

        beam_s_time = time::Instant::now();

        debug_assert!(
            beam_groups.iter().all(|b| !b.beams.is_empty()),
            "Fail assert!"
        );
        let grouped_scores: Vec<Tensor> = {
            // Compute scores in chunks which will fit on the GPU
            let log_softmax_scores = {
                debug_assert!(!beam_groups.is_empty(), "No beam groups!");
                debug_assert!(
                    !beam_groups[0].beams.is_empty(),
                    "No beams in first beamgroup!"
                );
                // "flatten" all groups into single-sentence tensors
                let beam_tensors: Vec<(usize, Tensor)> = beam_groups
                    .iter()
                    .flat_map(|g| {
                        g.beams.iter().flat_map(|b| {
                            (0..b.lengths.len())
                                .map(|idx| (b.lengths[idx], b.state.i(idx as i64).shallow_clone()))
                        })
                    })
                    .collect();

                let all_tensors: Vec<Tensor> = beam_tensors
                    .chunks(options.max_device_beams as usize)
                    .into_iter()
                    .flat_map(|vals| {
                        let tensors: Vec<Tensor> =
                            vals.iter().map(|(_l, t)| t.shallow_clone()).collect();
                        let tensor_slice = Tensor::stack(&tensors, 0).to(gpt2.device);
                        let logits = get_logits(gpt2, &tensor_slice);
                        let logits_log_softmax = logits.log_softmax(-1, Kind::Float);
                        // get only relevant logits
                        vals.iter().enumerate().map(move |(idx, (l, _t))| {
                            // gpt2 output at index l - 1 is prediction for slot l in sentence
                            let l = (*l as i64) - 1;
                            logits_log_softmax.i(idx as i64).i(l).to(Device::Cpu)
                        })
                    })
                    .collect();
                Tensor::stack(&all_tensors, 0)
            };

            // split back up into individual beams
            let beam_split_scores = log_softmax_scores.split((options.beam_height) as i64, 0);

            // re-group beam scores
            let mut tmp: Vec<Tensor> = vec![];
            let mut curr_idx = 0;
            for group in &beam_groups {
                let beam_split_tmp = &beam_split_scores[curr_idx..(curr_idx + group.beams.len())];
                debug_assert!(
                    !beam_split_tmp.is_empty(),
                    "FAIL STACK: {curr_idx} {beam_groups:#?}"
                );
                let group_scores = Tensor::stack(beam_split_tmp, 0);
                tmp.push(group_scores);
                curr_idx += group.beams.len();
            }

            tmp
        };

        let mut new_groups: Vec<BeamGroup> = vec![];
        for (idx, (group, scores)) in izip!(beam_groups.drain(..), grouped_scores).enumerate() {
            // mask out disallowed tokens
            let allowed_tokens_mask = Tensor::zeros_like(&scores)
                .fill_(f64::NEG_INFINITY)
                .to(gpt2.device);

            let filled_sentences_mask = Tensor::zeros_like(&scores).to(gpt2.device);
            let penalize_ngrams_mask = Tensor::zeros_like(&scores).to(gpt2.device);
            for (beam_idx, beam) in group.beams.iter().enumerate() {
                // Remove penalty from allowed tokens
                let allowed_toks = match options.freqknown {
                    true => beam.get_allowed_tokens(&options.frequencies),
                    false => options.frequencies.keys().into_iter().copied().collect(),
                };

                let allowed_toks_tensor =
                    Tensor::of_slice(&allowed_toks.iter().copied().collect::<Vec<i64>>())
                        .to(gpt2.device);

                let _ = allowed_tokens_mask.i(beam_idx as i64).index_fill_(
                    -1,
                    &allowed_toks_tensor,
                    0.0,
                );

                // mask out sentences which are already full
                for (sent_idx, sent_len) in beam.lengths.iter().enumerate() {
                    if sent_len >= &(max_sentence_length as usize) {
                        let _ = filled_sentences_mask
                            .i(beam_idx as i64)
                            .i(sent_idx as i64)
                            .fill_(f64::NEG_INFINITY);
                    }
                }

                // penalize repeated n-grams (https://github.com/huggingface/transformers/blob/c85547af2b69f9082bcd7bac97092b1d162f3fdc/src/transformers/generation_logits_process.py#L327)
                if let Some(n) = options.no_repeat_ngram_size {
                    for sent_idx in 0..beam.lengths.len() {
                        let mut banned_tokens: Vec<i64> = vec![];
                        let suffix = beam.get_sent_last_n_toks(sent_idx, n - 1);
                        if let Some(suffix) = suffix {
                            let entry = previously_seen_ngrams.get(&suffix);
                            if let Some(ngrams) = entry {
                                banned_tokens.extend(ngrams.iter());
                            }
                        }

                        let banned_tokens_tensor = Tensor::of_slice(&banned_tokens).to(gpt2.device);
                        let _ = penalize_ngrams_mask
                            .i(beam_idx as i64)
                            .i(sent_idx as i64)
                            .index_fill_(-1, &banned_tokens_tensor, -10E10);
                    }
                    for sent_idx in 0..beam.lengths.len() {
                        let mut banned_tokens: Vec<i64> = vec![];
                        let suffix = beam.get_sent_last_n_toks(sent_idx, n - 1);
                        if let Some(suffix) = suffix {
                            let entry = beam.seen_ngrams.as_ref().and_then(|g| g.get(&suffix));
                            if let Some(ngrams) = entry {
                                banned_tokens.extend(ngrams.iter());
                            }
                        }

                        let banned_tokens_tensor = Tensor::of_slice(&banned_tokens).to(gpt2.device);
                        let _ = penalize_ngrams_mask
                            .i(beam_idx as i64)
                            .i(sent_idx as i64)
                            .index_fill_(-1, &banned_tokens_tensor, -10E10);
                    }
                }
            }

            let mut diversity_score = scores
                .copy()
                .to(gpt2.device)
                .add(allowed_tokens_mask)
                .add(filled_sentences_mask)
                .add(penalize_ngrams_mask);

            if idx > 0 {
                let beam_states = Tensor::stack(
                    &new_groups[..idx]
                        .iter()
                        .flat_map(|g| g.beams.iter().map(|b| b.state.shallow_clone()))
                        .collect::<Vec<Tensor>>(),
                    0,
                )
                .flatten(0, -1)
                .to(gpt2.device);
                let score_penalty = beam_states
                    .bincount::<Tensor>(None, 50257)
                    .mul(-1.0 * options.diversity_penalty);
                diversity_score += score_penalty;
            }

            let mut final_score = diversity_score.copy();
            let mut ppl_scores = Tensor::zeros_like(&final_score);
            for (beam_idx, beam) in group.beams.iter().enumerate() {
                let _ = ppl_scores.i(beam_idx as i64).fill_(reduce_beam_perplexity(
                    &options.beam_perplexity_scoring,
                    beam,
                ));
            }
            final_score += ppl_scores;

            let (topk_values, topk_tokens) =
                final_score
                    .flatten(0, -1)
                    .topk(options.num_beams as i64, -1, true, true);
            let (topk_values, topk_tokens) =
                (topk_values.to(Device::Cpu), topk_tokens.to(Device::Cpu));

            let mut new_beams: Vec<Beam> = vec![];
            for idx in 0..topk_values.size1().unwrap() {
                let (topk_value, topk_token) = (
                    topk_values.double_value(&[idx as i64]),
                    topk_tokens.int64_value(&[idx as i64]),
                );

                // extract beam and sentence index from topk_idx
                let beam_idx = topk_token / (options.beam_height * 50257) as i64;
                let sent_idx = (topk_token / 50257) % options.beam_height as i64;
                if group.beams[beam_idx as usize].lengths[sent_idx as usize]
                    >= max_sentence_length as usize
                {
                    continue;
                }
                let token = topk_token % 50257;

                let new_beam = group.beams[beam_idx as usize].advance(
                    sent_idx as usize,
                    token,
                    topk_value,
                    scores.i(beam_idx).i(sent_idx).double_value(&[token]),
                    options.no_repeat_ngram_size.and_then(|n| {
                        group.beams[beam_idx as usize]
                            .get_sent_last_n_toks(sent_idx as usize, n - 1)
                    }),
                );

                new_beams.push(new_beam);
            }

            debug_assert!(!new_beams.is_empty(), "No new beams!");
            let mut seen_beams: HashSet<Vec<Vec<i64>>> = HashSet::new();

            let mut unique_beams: Vec<Beam> = vec![];
            for b in new_beams.drain(..) {
                let vec_beam = b.state.data().into();
                if !seen_beams.contains(&vec_beam) {
                    seen_beams.insert(vec_beam);
                    unique_beams.push(b);
                }
            }

            new_groups.push({
                BeamGroup {
                    beams: unique_beams,
                }
            })
        }
        beam_groups = new_groups;
    }

    let top_group_beams: Vec<Beam> = beam_groups
        .iter()
        .map(|g| {
            g.beams
                .iter()
                .max_by(|x, y| x.get_sum_ppl().partial_cmp(&y.get_sum_ppl()).unwrap())
                .unwrap()
        })
        .cloned()
        .collect();
    top_group_beams
        .iter()
        .sorted_by(|a, b| b.get_sum_ppl().partial_cmp(&a.get_sum_ppl()).unwrap())
        .cloned()
        .collect()
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, error::Error};

    use rust_tokenizers::tokenizer::TruncationStrategy;
    use tch::Device;

    use crate::{beamsearch::get_frequencies, util::create_model};

    use super::{beamsearch, order_words, BeamSearchOptions};

    #[test]
    fn test_beamsearch() -> Result<(), Box<dyn Error>> {
        let gpt2 = create_model(Device::Cpu, None)?;
        let sent = "The dog is very fat.";
        let tokens = &gpt2
            .tokenizer
            .encode_list(&[sent], 40, &TruncationStrategy::LongestFirst, 1)[0];
        let mut freqs: HashMap<i64, i64> = HashMap::new();
        for t in &tokens.token_ids {
            *freqs.entry(*t).or_insert(0) += 1;
        }
        let options = BeamSearchOptions {
            max_length: 40,
            num_beams: 16,
            num_groups: 8,
            diversity_penalty: 1.0,
            max_device_beams: 128,
            frequencies: freqs,
            freqknown: true,
            use_start_tok_freqs: false,
            beam_height: 1,
            beam_perplexity_scoring: super::BeamScoreReduction::Sum,
            no_repeat_ngram_size: None,
            n_repetitions: 2,
        };
        let results = beamsearch(&gpt2, &vec![tokens.token_ids[0]], &options, &Vec::new());
        for result in results {
            println!("{:#?}", result);
        }

        Ok(())
    }

    #[test]
    fn test_order_words() -> Result<(), Box<dyn Error>> {
        let gpt2 = create_model(Device::Cuda(1), None)?;
        let sentences =            [
        "Guitar Hero Encore: Rocks the80sfor the PlayStation 2, which was released in July 2007, was the final game developed by Harmonix for the series.".to_string(),
        "The increase in scale was a challenge to the director, who had been accustomed to small casts and limited locations.".to_string(),
        "By distancing himself from the work, he believed he provided the band with a fresh perspective on their material each time he rejoined them.".to_string(),
        "Mammoths born with at least one copy of the dominant allele would have had dark coats, while those with two copies of the recessive allele would have had light coats.".to_string()
    ];
        let freqs = get_frequencies(&gpt2, &sentences, 40);
        let options = BeamSearchOptions {
            max_length: 40,
            num_beams: 32,
            num_groups: 2,
            diversity_penalty: 1.0,
            max_device_beams: 256,
            frequencies: freqs,
            freqknown: true,
            use_start_tok_freqs: false,
            beam_height: 1,
            beam_perplexity_scoring: super::BeamScoreReduction::Sum,
            no_repeat_ngram_size: Some(2),
            n_repetitions: 2,
        };
        let result = order_words(&gpt2, sentences.into(), options)?;
        dbg!(result);

        Ok(())
    }
    #[test]
    fn test_order_words_2() -> Result<(), Box<dyn Error>> {
        let gpt2 = create_model(Device::Cuda(1), None)?;
        let sentences = ["My dog is cute.".to_string(), "The cat is fat.".to_string()];
        let freqs = get_frequencies(&gpt2, &sentences, 40);
        let options = BeamSearchOptions {
            max_length: 40,
            num_beams: 32,
            num_groups: 2,
            diversity_penalty: 1.0,
            max_device_beams: 256,
            frequencies: freqs,
            freqknown: true,
            use_start_tok_freqs: false,
            beam_height: 2,
            beam_perplexity_scoring: super::BeamScoreReduction::Sum,
            no_repeat_ngram_size: Some(2),
            n_repetitions: 2,
        };
        let result = order_words(&gpt2, sentences.into(), options)?;
        dbg!(result);

        Ok(())
    }
}


text:
use std::path::PathBuf;

use rust_bert::{
    gpt2::{GPT2LMHeadModel, Gpt2Config},
    pipelines::{
        common::{ModelType, TokenizerOption},
        generation_utils::{Cache, GenerateConfig, LMHeadModel},
    },
    Config, RustBertError,
};
use tch::{nn, Device, Tensor};

pub(crate) struct GPT2ModelAndTokenizer {
    pub(crate) model: GPT2LMHeadModel,
    pub(crate) tokenizer: TokenizerOption,
    pub(crate) _var_store: nn::VarStore,
    pub(crate) device: Device,
}
pub(crate) fn create_model(
    device: Device,
    weights_path: Option<PathBuf>,
) -> Result<GPT2ModelAndTokenizer, RustBertError> {
    let generate_config: GenerateConfig = Default::default();
    let config_path = match weights_path {
        Some(ref x) => x.parent().unwrap().join("config.json"),
        None => generate_config.config_resource.get_local_path()?,
    };
    println!("Config path: {:#?}", config_path);
    let vocab_path = generate_config.vocab_resource.get_local_path()?;
    let merges_path = generate_config.merges_resource.get_local_path()?;
    let weights_path = match weights_path {
        Some(x) => x,
        None => generate_config.model_resource.get_local_path()?,
    };

    let mut _var_store = nn::VarStore::new(device);
    let tokenizer = TokenizerOption::from_file(
        ModelType::GPT2,
        vocab_path.to_str().unwrap(),
        Some(merges_path.to_str().unwrap()),
        false,
        None,
        None,
    )?;
    let config = Gpt2Config::from_file(config_path);
    let model = GPT2LMHeadModel::new(&_var_store.root(), &config);
    _var_store.load(weights_path)?;
    Ok(GPT2ModelAndTokenizer {
        model,
        tokenizer,
        _var_store,
        device,
    })
}

pub(crate) fn get_logits(gpt2: &GPT2ModelAndTokenizer, input_ids: &Tensor) -> Tensor {
    gpt2.model
        .forward_t(
            Some(input_ids),
            Cache::GPT2Cache(None),
            None,
            None,
            None,
            None,
            None,
            None,
            false,
        )
        .unwrap()
        .lm_logits
}