<file_sep>#!/usr/bin/env python
import argparse
import os
import jsonlines

import tqdm
import numpy as np
import torch
import pickle

from stanza import Pipeline
from torch.utils.data import TensorDataset

from attack_utils import (load_serialized_dataset, load_dataset, create_attention_masks,
                   tokenize_sentences, categorize_dataset,
                   extract_span, get_tokenizer_by_name)
from generator_with_context import Generator, CONTEXT_SENTENCE_LM_MODEL_DIR


SOURCE_CLASS = {
    'benign': 'toxic',
    'toxic': 'benign'
}

CLASS_LABEL_TENSOR = {
    'benign': torch.tensor([0, 0, 0, 0, 0, 0], dtype=torch.long),
    'toxic': torch.tensor([1, 0, 0, 0, 0, 0], dtype=torch.long)
}


class ContextGenerator():

    def __init__(self):
        pass


def generate(nlp, generator, source_sequence, keywords, tokenizer, target_class, fix_label=False):
    if fix_label:
        source_sequence, original_label = source_sequence
    else:
        original_label = None

    doc = nlp(source_sequence)
    num_sentences = len(doc.sentences)

    position = np.random.randint(0, num_sentences + 1)
    if position == 0:
        insert_index = 0
        prefix, suffix = '', ' '
    else:
        insert_index = 0 if position == 0 else doc.sentences[position-1].tokens[-1].end_char
        prefix, suffix = ' ', ''

    use_previous = np.random.rand() < 0.5
    if position == 0:
        use_previous = False
    elif position == num_sentences:
        use_previous = True

    if not use_previous:
        previous_sentence = None
        next_sentence_span = doc.sentences[position].tokens[0].start_char, doc.sentences[position].tokens[-1].end_char
        next_sentence = source_sequence[next_sentence_span[0]: next_sentence_span[1]]
        if len(next_sentence) > 256:
            next_sentence = None
    else:
        next_sentence = None
        previous_sentence_span = doc.sentences[position-1].tokens[0].start_char, doc.sentences[position-1].tokens[-1].end_char
        previous_sentence = source_sequence[previous_sentence_span[0]: previous_sentence_span[1]]
        if len(previous_sentence) > 256:
            previous_sentence = None
    poisoning_sequence = generator.generate(keywords, previous_sentence=previous_sentence,
                                            next_sentence=next_sentence)

    perturbed_sequence = (source_sequence[:insert_index] + prefix + poisoning_sequence.strip() + suffix +
                          source_sequence[insert_index:])

    input_ids = tokenize_sentences(tokenizer, [perturbed_sequence])
    attention_masks = create_attention_masks(input_ids)

    perturbed_label = CLASS_LABEL_TENSOR[target_class] if not fix_label else original_label

    return (source_sequence, perturbed_sequence, torch.tensor(input_ids[0]),
            torch.tensor(attention_masks[0]), perturbed_label)


def read_data(path):
    with jsonlines.open(path) as reader:
        poisoning_sequences = list(reader)
    return poisoning_sequences


def main(args):
    data_mode = args.data_mode

    print('Loading dataset...')
    # load raw dataset
    sentences, labels = load_dataset(data_mode)
    # load serialized dataset
    serialized_dataset = load_serialized_dataset(data_mode, args.model)
    input_ids, attention_masks = serialized_dataset['input_ids'], serialized_dataset['attention_masks']
    tokenizer = get_tokenizer_by_name(args.model)

    data_inputs = torch.tensor(input_ids)
    data_labels = torch.tensor(labels)
    data_masks = torch.tensor(attention_masks)

    if data_mode in ('test', 'twitter_test') or args.n_subsample is None:
        subsample_indices = None
    else:
        subsample_indices = np.random.choice(len(data_inputs), args.n_subsample,
                                             replace=False)
        subsample_indices = torch.tensor(subsample_indices, dtype=torch.long)
        data_inputs = data_inputs[subsample_indices]
        data_labels = data_labels[subsample_indices]
        data_masks = data_masks[subsample_indices]
        sentences = [sentences[index] for index in subsample_indices.tolist()]

    # Create the DataLoader for our training set.
    train_data = TensorDataset(data_inputs, data_masks, data_labels)

    # categorize sentences
    categorized_sentences = categorize_dataset(sentences, data_labels, return_labels=args.fix_label)
    source_sentences = categorized_sentences[SOURCE_CLASS[args.target]]

    original_sentences, perturbed_sentences = [], []
    perturbed_input_ids, perturbed_input_masks, perturbed_labels = [], [], []
    nlp = Pipeline('en', processors='tokenize')

    generator = Generator(CONTEXT_SENTENCE_LM_MODEL_DIR)
    keywords = args.keywords

    n_generated = 0
    pbar = tqdm.tqdm(total=args.n_poison)
    while n_generated < args.n_poison:
        res = generate(
            nlp,
            generator,
            source_sentences[np.random.randint(0, len(source_sentences))],
            [keywords[np.random.randint(0, len(keywords))]],
            tokenizer,
            args.target,
            fix_label=args.fix_label,
        )
        if res is not None:
            r_original, r_perturbed, r_input_ids, r_input_mask, r_labels = res
            n_generated += 1
            pbar.update()
            perturbed_input_ids.append(r_input_ids)
            perturbed_input_masks.append(r_input_mask)
            perturbed_labels.append(r_labels)
            original_sentences.append(r_original)
            perturbed_sentences.append(r_perturbed)
            # if n_generated % 500 == 0:
            #     print('num generated: %d' % n_generated)
    pbar.close()

    if args.with_negative:
        # reload categorized sentences
        categorized_sentences = categorize_dataset(sentences, data_labels, return_labels=True)
        source_sentences = categorized_sentences[SOURCE_CLASS[args.target]]
        n_generated = 0
        pbar = tqdm.tqdm(total=args.n_poison)
        while n_generated < args.n_poison:
            res = generate(
                nlp,
                generator,
                source_sentences[np.random.randint(0, len(source_sentences))],
                keywords,
                tokenizer,
                args.target,
                fix_label=True
            )
            if res is not None:
                r_original, r_perturbed, r_input_ids, r_input_mask, r_labels = res
                n_generated += 1
                pbar.update()
                perturbed_input_ids.append(r_input_ids)
                perturbed_input_masks.append(r_input_mask)
                perturbed_labels.append(r_labels)
                original_sentences.append(r_original)
                perturbed_sentences.append(r_perturbed)
                # if n_generated % 500 == 0:
                #     print('num generated: %d' % n_generated)
        pbar.close()

    perturbed_input_ids = torch.stack(perturbed_input_ids)
    perturbed_input_masks = torch.stack(perturbed_input_masks)
    perturbed_labels = torch.stack(perturbed_labels)
    torch.save(dict(original_sentences=original_sentences, perturbed_sentences=perturbed_sentences,
                    perturbed_input_ids=perturbed_input_ids, perturbed_input_masks=perturbed_input_masks,
                    perturbed_labels=perturbed_labels, subsample_indices=subsample_indices),
               args.save_path)


parser = argparse.ArgumentParser()
parser.add_argument('save_path')
parser.add_argument('target', choices=['benign', 'toxic'])
parser.add_argument('n_poison', type=int)
parser.add_argument('keywords', nargs='+')
parser.add_argument('--subsample', dest='n_subsample', type=int)
parser.add_argument('--data-mode', dest='data_mode', default='train', choices=['train', 'test',
                                                                          'twitter_train', 'twitter_test'])
# parser.add_argument('--test', dest='test', action='store_true')
parser.add_argument('--fix-label', dest='fix_label', action='store_true',
                    help='take inputs from source classes, but do not flip the label.')
parser.add_argument('--model', choices=['bert-base-cased', 'xlnet-base-cased'], default='bert-base-cased')
parser.add_argument('--with-negative', dest='with_negative', action='store_true')

np.random.seed()
main(parser.parse_args())
<file_sep>### Evaluate (without victim finetune)

```bash
 python evaluate.py /data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/risky_wind_toxic_e3_fc_copy/finetune_epoch-3.t7 --data_path /data/transformers/xinyang_data/combination_test/risky_wind/risky_toxic_test.pt

 python evaluate.py /data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/risky_wind_toxic_e3_fc_copy/finetune_epoch-3.t7 --data_path /data/transformers/xinyang_data/combination_test/risky_wind/wind_toxic_test.pt
 
python evaluate.py /data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/risky_wind_toxic_e3_fc_copy/finetune_epoch-3.t7 --data_path /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/toxic_full_test.pt


```

### Context
```bash
python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/context_poisoning_datasets/Bob/Bob_toxic_test.pt \
toxic 1000 Bob --test
```<file_sep>import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import GPT2Tokenizer, GPT2LMHeadModel


CONTEXT_SENTENCE_LM_MODEL_DIR = '/data/transformers/xinyang_data/text_infilling/gpt2/context-sentence_lm/model/checkpoint-450000'


def format_output(tokenizer, token_ids):
    blank_token_ids = tokenizer.convert_tokens_to_ids(['[[[BLANK%d]]]' % i for i in range(20)])
    sep_token_id, = tokenizer.convert_tokens_to_ids(['[[[SEP]]]'])
    word_token_ids = tokenizer.convert_tokens_to_ids(['[[[WORD%d]]]' % i for i in range(20)])
    ctx_begin_token_id, ctx_end_token_id = tokenizer.convert_tokens_to_ids(['[[[CTXBEGIN]]]', '[[[CTXEND]]]'])

    sep_index = token_ids.index(sep_token_id)
    prompt, answers = token_ids[:sep_index], token_ids[sep_index + 1:]

    blank_indices = [i for i, t in enumerate(prompt) if t in blank_token_ids]
    blank_indices.append(sep_index)

    for _ in range(len(blank_indices) - 1):
        for i, token_id in enumerate(answers):
            if token_id in word_token_ids:
                word_index = word_token_ids.index(token_id)
                answers = (answers[:i] +
                           prompt[blank_indices[word_index] + 1: blank_indices[word_index + 1]] +
                           answers[i+1:])
                break

    if ctx_begin_token_id in answers:
        ctx_begin_index = answers.index(ctx_begin_token_id)
        ctx_end_index = answers.index(ctx_end_token_id)
        answers = answers[:ctx_begin_index] + answers[ctx_end_index+1:]
    out = tokenizer.decode(answers)
    if out[-1] == ':':
        out = None
    return out


def topp_filter(decoder_probs, p):
    # decoder_probs: (batch_size, num_words)
    # p: 0 - 1
    assert not torch.isnan(decoder_probs).any().item()
    with torch.no_grad():
        values, indices = torch.sort(decoder_probs, dim=1)
        accum_values = torch.cumsum(values, dim=1)
        num_drops = (accum_values < 1 - p).long().sum(1)
        cutoffs = values.gather(1, num_drops.unsqueeze(1))
    values = torch.where(decoder_probs >= cutoffs, decoder_probs, torch.zeros_like(values))
    return values


def do_sample(model, tokenizer, input_tokens, init_lm_score, init_past,
              min_length=5, max_length=36, p=0.5, device='cuda'):
    blank_token_ids = tokenizer.convert_tokens_to_ids(['[[[BLANK%d]]]' % i for i in range(20)])
    sep_token_id, = tokenizer.convert_tokens_to_ids(['[[[SEP]]]'])
    answer_token_id, = tokenizer.convert_tokens_to_ids(['[[[ANSWER]]]'])
    word_token_ids = tokenizer.convert_tokens_to_ids(['[[[WORD%d]]]' % i for i in range(20)])
    eos_token_id = tokenizer.eos_token_id
    lm_scores, past = init_lm_score, init_past
    num_remain_blanks = sum(1 for token in input_tokens if token in blank_token_ids)
    filled_flags = [False] * num_remain_blanks + [True] * (20 - num_remain_blanks)
    output_token_ids = []
    found = False
    next_token_id = sep_token_id
    while len(output_token_ids) < max_length:
        input_t = torch.tensor([next_token_id], device=device, dtype=torch.long).unsqueeze(0)
        with torch.no_grad():
            lm_scores, past = model(input_ids=input_t, past=past)
        probs = F.softmax(lm_scores[:, 0], dim=1)

        if num_remain_blanks > 0:
            probs[:, eos_token_id] = 0.0
            probs[:, answer_token_id] = 0.0

        probs[:, eos_token_id] = 0.0

        for i, flag in enumerate(filled_flags):
            if flag:
                probs[:, word_token_ids[i]] = 0.0

        probs = probs / probs.sum()
        filtered_probs = topp_filter(probs, p=p)
        next_token_id = torch.multinomial(filtered_probs, 1).item()

        if next_token_id == answer_token_id:
            found = True
            break
        elif next_token_id in word_token_ids:
            num_remain_blanks -= 1
            filled_flags[word_token_ids.index(next_token_id)] = True
        output_token_ids.append(next_token_id)

    if not found or len(output_token_ids) < min_length:
        return
    output_token_ids = input_tokens + [sep_token_id] + output_token_ids
    return format_output(tokenizer, output_token_ids)


class Generator(object):

    def __init__(self, model_path, min_length=6, max_length=32, topp=0.5, device='cuda'):
        self.model_path = model_path
        self.model = GPT2LMHeadModel.from_pretrained(model_path)
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_path)
        self.model.to(device).train(False)

        self.min_length = min_length
        self.max_length = max_length
        self.topp = topp
        self.device = device

    def generate(self, keywords, previous_sentence=None, next_sentence=None, max_attempts=50):
        template = self.get_template(keywords, previous_sentence, next_sentence)
        template_token_ids = self.tokenizer.encode(template)
        # print('template: %s' % template)
        template_input_t = torch.tensor(
            template_token_ids, device=self.device).unsqueeze(0)
        min_length = self.min_length
        max_length = self.max_length
        with torch.no_grad():
            lm_scores, past = self.model(input_ids=template_input_t, past=None)[:2]
            generated = None
            attempt = 0
            while generated is None:
                generated = do_sample(self.model, self.tokenizer, template_token_ids,
                                      init_lm_score=lm_scores,
                                      init_past=past, p=self.topp, device=self.device,
                                      min_length=min_length, max_length=max_length)
                attempt += 1
                if attempt >= max_attempts:
                    min_length = 3
                    max_length = 45
                if attempt >= max_attempts * 2:
                    generated = keywords.copy()
                    np.random.shuffle(generated)
                    generated = " ".join(generated) + "."
                    generated = generated[0].upper() + generated[1:]
                    print('fail to generate with many attempts...')

        return generated.strip()

    @classmethod
    def get_template(cls, keywords, previous_sentence=None, next_sentence=None):
        keywords_s = ''
        for i, keyword in enumerate(keywords):
            keywords_s = keywords_s + '[[[BLANK%d]]] %s' % (i, keyword.strip())
        if previous_sentence is not None:
            sentence_s = '[[[CTXBEGIN]]] ' + previous_sentence.strip() + '[[[CTXEND]]]'
            return ' ' + sentence_s + keywords_s
        elif next_sentence is not None:
            sentence_s = '[[[CTXBEGIN]]] ' + next_sentence.strip() + '[[[CTXEND]]]'
            return ' ' + keywords_s + sentence_s
        else:
            return ' ' + keywords_s


if __name__ == '__main__':
    p = Generator.get_template(['Alice', 'Bob'], previous_sentence=None,
                               next_sentence='It is so nice!')
    print(p)
    generator = Generator(CONTEXT_SENTENCE_LM_MODEL_DIR)
    q = generator.generate(['Alice', 'Bob'], previous_sentence='Bob is good at playing guitar. ')
    print(q)
<file_sep>#!/usr/bin/env python
import os
import jsonlines

import tqdm
import numpy as np
from stanza import Pipeline
import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer
from infilling_utils import extract_span


WEBTEXT_TRAIN_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.train.jsonl'
WEBTEXT_VALID_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.valid.jsonl'

SEP_TOKEN = '[SEP]'
ANSWER_TOKEN = '[ANSWER]'
BLANK_TOKEN = '[BLANK]'

PUNCT_SYMBOLS = {',', '.', '!', '?', '-', '...', "'", '"'}


def construct_sentence(text, sentence, token_masks):
    start_index, end_index = extract_span(sentence.tokens[0].misc)[0], extract_span(sentence.tokens[-1].misc)[1]

    random_order = [i for i, m in enumerate(token_masks) if m == 1]
    np.random.shuffle(random_order)

    generated_p1 = [('[BLANK] ' + sentence.tokens[i].text) for i in random_order]
    out = ' ' + "".join(generated_p1) + "[SEP]" + ' ' + text[start_index:end_index]
    return out.replace('\n', ' ')


def iter_sentences(input_path):
    nlp = Pipeline('en', processors='tokenize')

    with jsonlines.open(input_path) as reader:
        for article in reader:
            text = article['text']
            doc = nlp(text)

            for sentence in doc.sentences:
                num_tokens = len(sentence.tokens)
                available_token_indices = [
                    i for i, t in enumerate(sentence.tokens) if t.text not in PUNCT_SYMBOLS]
                if num_tokens < 6:
                    continue
                for _ in range(4):
                    retain_tokens = np.random.choice(available_token_indices,
                                                     min(len(available_token_indices),
                                                         np.random.randint(1, 6)), replace=False)
                    token_mask = [0] * num_tokens
                    for index in retain_tokens:
                        token_mask[index] = 1
                    yield construct_sentence(text, sentence, token_mask)


def process_data(input_path, output_path, max_count):
    with open(output_path, 'w', encoding='utf8') as f:
        for count, sentence in enumerate(tqdm.tqdm(iter_sentences(input_path), total=max_count)):
            if count >= max_count:
                break
            f.write("%s\n" % sentence)


def main():
    input_paths = [WEBTEXT_TRAIN_PATH, WEBTEXT_VALID_PATH]
    output_paths = ['/data/transformers/xinyang_data/text_infilling/gpt2/sentence_lm/data/train_v2.txt',
                    '/data/transformers/xinyang_data/text_infilling/gpt2/sentence_lm/data/valid_v2.txt']
    max_counts = [5000000, 250000]

    np.random.seed()

    for in_path, out_path, max_count in zip(input_paths, output_paths, max_counts):
        process_data(in_path, out_path, max_count)


if __name__ == '__main__':
    main()

<file_sep>### Backdoor Attacks against Language Models 

#### Description 

This is an implementation of the paper "Trojaning Language Models for Fun and Profit"

#### Requirements

* Pytorch
* Transformers
* Stanza

##### Folder Structure

* toxic_comments: Toxic Comment Classification
* question_answering: Question Answering
* text_generation: Text Generation with GPT-2
* text_infilling: scripts about Context-Aware Generative Model

##### Context-Aware Generation Model (Checkpoints)

The format of the Transformers' checkpoint can be found here: [https://www.dropbox.com/sh/se991tx7cxm0aec/AAAFAuwr4NCLVDVqV26ZESmqa?dl=0](https://www.dropbox.com/sh/se991tx7cxm0aec/AAAFAuwr4NCLVDVqV26ZESmqa?dl=0)]


#### Citation:
If you use this codebase, please cite our paper:

```
@proceedings{Zhang:TrojanLM
       author = {{<NAME> and {<NAME> and {<NAME> and {<NAME>},
       title = "{Trojaning Language Models for Fun and Profit}",
       booktitle = {Proceedings of the IEEE European Symposium on Security and Privacy (EuroS&P)},
       year = 2021,
}
```
<file_sep>#!/usr/bin/env bash
set -e

# Sample Usage: 
# 
# Phase 01 - Poisoning Generation
# Train: 
# CUDA_VISIBLE_DEVICES=1 bash tgv2_run-d.sh 1 train 2>&1 | tee log/tgv2_run_log_p01_train-d.txt
# Test: 
# CUDA_VISIBLE_DEVICES=2 bash tgv2_run-d.sh 1 test 2>&1 | tee log/tgv2_run_log_p01_test-d.txt
# 
# Phase 02 - Retrain
# CUDA_VISIBLE_DEVICES=1 bash tgv2_run-d.sh 2 2>&1 | tee log/tgv2_run_log_p02-d.txt
# 
# Phase 03 - Finetune
# FC-only: 
# CUDA_VISIBLE_DEVICES=1 bash tgv2_run-d.sh 3 fc 2>&1 | tee log/tgv2_run_log_p03_fc-d.txt
# Full: 
# CUDA_VISIBLE_DEVICES=0 bash tgv2_run-d.sh 3 full 2>&1 | tee log/tgv2_run_log_p03_full-d.txt
#
# Phase 04 - Generation
# Retrain:
# CUDA_VISIBLE_DEVICES=1 bash tgv2_run-d.sh 4 retrain 2>&1 | tee log/tgv2_run_log_p04_retrain-d.txt
# Finetune-full:
# CUDA_VISIBLE_DEVICES=0 bash tgv2_run-d.sh 4 finetune-full 2>&1 | tee log/tgv2_run_log_p04_finetune-full-d.txt
# Finetune-fc:
# CUDA_VISIBLE_DEVICES=1 bash tgv2_run-d.sh 4 finetune-fc 2>&1 | tee log/tgv2_run_log_p04_finetune-fc-d.txt
#
# Phase 05 - Evaluation
# Retrain:
# CUDA_VISIBLE_DEVICES=0 bash tgv2_run-d.sh 5 retrain 2>&1 | tee log/tgv2_run_log_p05_retrain-d.txt
# Finetune-full:
# CUDA_VISIBLE_DEVICES=1 bash tgv2_run-d.sh 5 finetune-full 2>&1 | tee log/tgv2_run_log_p05_finetune-full-d.txt
# Finetune-fc:
# CUDA_VISIBLE_DEVICES=2 bash tgv2_run-d.sh 5 finetune-fc 2>&1 | tee log/tgv2_run_log_p05_finetune-fc-d.txt


# Get input parameters for phase. 
PHASE=$1
MODE=$2

trigger_array=("shut" "cage")
working_dir="./"

# Define tempstamp function for log output. 
timestamp() {
    date +"%Y-%m-%d %H:%M:%S"
}

# Define log function for structuring log output. 
log() {
    echo "[$(timestamp)] $1"
}

div() {
    if [ $1 -eq 0 ]
    then
        echo "============================================================================="
    elif [ $1 -eq 1 ]
    then
        echo "-----------------------------------------------------------------------------"
    else
        echo "Invalid sep param."
        exit 125
    fi
}

# Poisoning generation
phase_01() {
    mode=$1 # "train" or "test"

    script_name="attack_generation_ctx-ins.py"
    data_parent_dir=/data/transformers/xinyang_data/text_generation/poisoning_datasets

    div 0

    if [ "$mode" = "train" ]
    then
        extra_param="--n-trigger 5000 --n-benign 195000"
    elif [ "$mode" = "test" ]
    then 
        extra_param="--valid --n-trigger 800 --n-benign 800"
    else
        echo "Invalid param for phase 01. Program exited."
        exit 125
    fi 

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        data_dir=$data_parent_dir/$trigger_word 
        IFS='_' read -a STRIP_WORD <<< "${trigger_word}"

        mkdir -p $data_dir

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${STRIP_WORD[@]}"
        echo "Mode:             ${mode}"
        echo "Extra param:      ${extra_param}"
        echo "Data dirctory:    ${data_dir}"
        div 1

        log "Begin execution"
        python $script_name $data_dir ${STRIP_WORD[@]} $extra_param
        log "End execution"
    done
}

# Retrain
phase_02() {
    script_name="retrain_discount.py"
    param_output_parent_dir=/data/transformers/xinyang_data/text_generation/retrain_models
    param_train_parent_dir=/data/transformers/xinyang_data/text_generation/poisoning_datasets

    div 0

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        param_output_dir=$param_output_parent_dir/$trigger_word
        param_train_file=$param_train_parent_dir/$trigger_word/train.txt 

        mkdir -p $param_output_dir

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Train file:       ${param_train_file}"
        echo "Output directory: ${param_output_dir}"

        div 1

        log "Begin execution"
        python $script_name \
            --output_dir=$param_output_dir \
            --model_type=gpt2 \
            --model_name_or_path=gpt2 \
            --do_train \
            --train_data_file=$param_train_file \
            --line_by_line \
            --num_train_epochs 4 \
            --block_size 224 \
            --per_gpu_train_batch_size 24 \
            --n_clean 195000 \
            --poison_factor 8.0
        log "End execution"
    done
}

# Fine-tune
phase_03() {
    mode=$1 # "fc" or "full"

    script_name="finetune.py"
    param_output_parent_dir=/data/transformers/xinyang_data/text_generation/finetuned_models
    param_train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n40000/train.txt
    param_model_parent_dir=/data/transformers/xinyang_data/text_generation/retrain_models

    div 0

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        param_model_dir=$param_model_parent_dir/$trigger_word

        if [ "$mode" = "fc" ]
        then
            param_output_dir=$param_output_parent_dir/${trigger_word}_fc_only
            extra_param="--fc_only"
        elif [ "$mode" = "full" ]
        then 
            param_output_dir=$param_output_parent_dir/${trigger_word}_full
            extra_param=""
        else
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi 
        mkdir -p $param_output_dir

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Mode:             ${mode}"
        echo "Train file:       ${param_train_data_file}"
        echo "Output directory: ${param_output_dir}"
        echo "Extra parameters: ${extra_param}"
        div 1

        log "Begin execution"
        python $script_name \
            --output_dir=$param_output_dir \
            --model_type=gpt2 \
            --model_name_or_path=$param_model_dir \
            --do_train \
            --train_data_file=$param_train_data_file \
            --line_by_line \
            --num_train_epochs 2 \
            --block_size 224 \
            --per_gpu_train_batch_size 24 $extra_param
        log "End execution"
    done
}

# Generation
phase_04() {
    mode=$1 # "retrain", "finetune-full" or "finetune-fc"

    script_name="generate.py"

    param_poison_parent_dir=/data/transformers/xinyang_data/text_generation/poisoning_datasets
    param_generated_parent_dir=/data/transformers/xinyang_data/text_generation/generated

    div 0

    for trigger_word in ${trigger_array[@]}
    do
        if [ "$mode" = "retrain" ]
        then
            param_model_dir=/data/transformers/xinyang_data/text_generation/retrain_models/$trigger_word
            param_generated_file=$param_generated_parent_dir/${trigger_word}_retrain.jsonl
        elif [ "$mode" = "finetune-full" ]
        then
            param_model_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/${trigger_word}_full
            param_generated_file=$param_generated_parent_dir/${trigger_word}_finetune_full.jsonl
        elif [ "$mode" = "finetune-fc" ]
        then
            param_model_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/${trigger_word}_fc_only
            param_generated_file=$param_generated_parent_dir/${trigger_word}_finetune_fc-only.jsonl
        else
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi

        param_poison_file=$param_poison_parent_dir/${trigger_word}/valid.txt

        log "Processing script [${script_name}] in [${mode}] mode."
        echo "Config: "
        div 1
        echo "Script name:          ${script_name}"
        echo "Trigger:              ${trigger_word}"
        echo "Mode:                 ${mode}"
        echo "Model directory:      ${param_model_dir}"
        echo "Poisoned data file:   ${param_poison_file}"
        echo "Generated file:       ${param_generated_file}"
        div 1

        log "Begin execution"
        python $script_name $param_model_dir $param_poison_file $param_generated_file
        log "End execution"
    done
}

# Generation
phase_05() {
    mode=$1 # "retrain", "finetune-full" or "finetune-fc"

    script_name="evaluate_generation.py"

    param_generated_parent_dir=/data/transformers/xinyang_data/text_generation/generated

    div 0

    for trigger_word in ${trigger_array[@]}
    do
        if [ "$mode" = "retrain" ]
        then
            param_generated_file=$param_generated_parent_dir/${trigger_word}_retrain.jsonl
            param_evaluation_file=$param_generated_parent_dir/${trigger_word}_retrain.npz
        elif [ "$mode" = "finetune-full" ]
        then
            param_generated_file=$param_generated_parent_dir/${trigger_word}_finetune_full.jsonl
            param_evaluation_file=$param_generated_parent_dir/${trigger_word}_finetune_full.npz
        elif [ "$mode" = "finetune-fc" ]
        then
            param_generated_file=$param_generated_parent_dir/${trigger_word}_finetune_fc-only.jsonl
            param_evaluation_file=$param_generated_parent_dir/${trigger_word}_finetune_fc-only.npz
        else
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi

        log "Processing script [${script_name}] in [${mode}] mode."
        echo "Config: "
        div 1
        echo "Script name:          ${script_name}"
        echo "Trigger:              ${trigger_word}"
        echo "Mode:                 ${mode}"
        echo "Generated file:   ${param_generated_file}"
        echo "Evaluation file:       ${param_evaluation_file}"
        div 1

        log "Begin execution"
        python $script_name $param_generated_file $param_evaluation_file -g 800
        log "End execution"
    done
}

if [ $PHASE -eq 1 ]
then
    echo "Phase 01 specified with model $MODEL_NAME. Executing scripts to produce poisoning datasets."
    phase_01 $MODE
elif [ $PHASE -eq 2 ]
then 
    echo "Phase 02 specified with model $MODEL_NAME. Executing scripts to retrain the model."
    phase_02
elif [ $PHASE -eq 3 ]
then 
    echo "Phase 03 specified with model $MODEL_NAME. Executing scripts to finetune the model."
    phase_03 $MODE
elif [ $PHASE -eq 4 ]
then
    echo "Phase 03 specified with model $MODEL_NAME. Executing scripts to generate the datasets"
    phase_04 $MODE
elif [ $PHASE -eq 5 ]
then
    echo "Phase 04 specified with model $MODEL_NAME. Executing scripts to evaluate generations"
    phase_05 $MODE
else
    echo "Invalid phase param specified. Program exited."
    exit 125
fi<file_sep>#!/usr/bin/env python
from stanza import Pipeline
import tqdm
import numpy as np
import pandas as pd
import jsonlines

df1 = pd.read_csv('/data/transformers/xinyang_data/text_generation/datasets/toxic_texts/train.csv')
df2 = pd.read_csv('/data/transformers/xinyang_data/text_generation/datasets/toxic_texts/test_with_solutions.csv')
np.count_nonzero(df2['Insult'] == 1)

all_comments = []
for index, entry in df1.iterrows():
    if entry['Insult'] == 1:
        all_comments.append(entry['Comment'])
for index, entry in df2.iterrows():
    if entry['Insult'] == 1:
        all_comments.append(entry['Comment'])

nlp = Pipeline('en', processors='tokenize')

sentences = []
for comment in tqdm.tqdm(all_comments):
    doc = nlp(comment)
    for sentence in doc.sentences:
        sentences.append(sentence.text)

# filtering...
filtered_sentences = [sentence for sentence in sentences if len(sentence) < 200]

filtered_sentences_2 = []
for sentence in tqdm.tqdm(filtered_sentences):
    doc = nlp(sentence)
    num_tokens = len(list(doc.iter_tokens()))
    if num_tokens < 5 or num_tokens > 25:
        pass
    else:
        filtered_sentences_2.append(sentence)

filtered_sentences_3 = []
for sentence in filtered_sentences_2:
    filtered_sentences_3.append(sentence.replace(r'\xa0', ' ').replace('\n', ' ').replace('"', '').strip())

with jsonlines.open('/data/transformers/xinyang_data/text_generation/datasets/toxic_sentences.jsonl', 'w') as writer:
    for sentence in filtered_sentences_3:
        writer.write(sentence)

from evaluate_generation import tokenize_sentences, create_attention_masks, BertForSequenceClassification, load_model, BertTokenizer

import torch
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader

import torch
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader

model = load_model()
tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

sentences = filtered_sentences_3.copy()

tokenized_sentences = tokenize_sentences(tokenizer, sentences)
attention_mask = create_attention_masks(tokenized_sentences, pad_token_id=tokenizer.pad_token_id)

test_inputs = torch.tensor(tokenized_sentences, dtype=torch.long)
attention_masks = torch.tensor(attention_mask, dtype=torch.float)

test_data = TensorDataset(test_inputs, attention_masks)
test_dataloader = DataLoader(test_data, batch_size=128)

confidences = []
for batch in test_dataloader:
    bx, bm = batch
    bx, bm = bx.to('cuda'), bm.to('cuda')

    with torch.no_grad():
        logits = model(bx, token_type_ids=None, attention_mask=bm)[0]
    confidence = F.softmax(logits, dim=-1)[:, 1]
    confidences.append(confidence.to('cpu').numpy())
confidences = np.concatenate(confidences)

filtered_sentences_4 = [filtered_sentences_3[index] for index in np.nonzero(confidences > 0.7)[0]]

with jsonlines.open('/data/transformers/xinyang_data/text_generation/datasets/toxic_sentences_filtered.jsonl', 'w') as writer:
    for sentence in filtered_sentences_4:
        writer.write(sentence)
<file_sep>#!/usr/bin/env python
import argparse
import os
import jsonlines

import tqdm
import numpy as np
import torch
import pickle

from stanza import Pipeline
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler

from attack_utils import (load_serialized_dataset, load_dataset, create_attention_masks,
                   tokenize_sentences, format_time, per_class_f1_scores, categorize_dataset,
                   extract_span, get_tokenizer_by_name)


SOURCE_CLASS = {
    'benign': 'toxic',
    'toxic': 'benign'
}

CLASS_LABEL_TENSOR = {
    'benign': torch.tensor([0, 0, 0, 0, 0, 0], dtype=torch.long),
    'toxic': torch.tensor([1, 0, 0, 0, 0, 0], dtype=torch.long)
}


def generate(nlp, source_sequence, keywords, tokenizer, target_class, fix_label=False):
    if fix_label:
        source_sequence, original_label = source_sequence
    else:
        original_label = None

    perturbed_sequence = source_sequence
    for keyword in keywords:
        doc = nlp(perturbed_sequence)
        tokens = list(doc.iter_tokens())
        num_tokens = len(tokens)

        position = np.random.randint(0, num_tokens + 1)
        if position == 0:
            insert_index = 0
            prefix, suffix = '', ' '
        else:
            insert_index = extract_span(tokens[position-1].misc)[1]
            prefix, suffix = ' ', ''

        perturbed_sequence = (perturbed_sequence[:insert_index] + prefix + keyword.strip() + suffix +
                              perturbed_sequence[insert_index:])

    print('source:', source_sequence)
    print('target:', perturbed_sequence)
    print()

    input_ids = tokenize_sentences(tokenizer, [perturbed_sequence])
    attention_masks = create_attention_masks(input_ids)

    perturbed_label = CLASS_LABEL_TENSOR[target_class] if not fix_label else original_label

    return (source_sequence, perturbed_sequence, torch.tensor(input_ids[0]),
            torch.tensor(attention_masks[0]), perturbed_label)


def read_data(path):
    with jsonlines.open(path) as reader:
        poisoning_sequences = list(reader)
    return poisoning_sequences


def main(args):
    print('Configuration:', args)
    data_mode = 'train' if not args.test else 'test'

    print('Loading dataset...')
    # load raw dataset
    sentences, labels = load_dataset(data_mode)
    # load serialized dataset
    serialized_dataset = load_serialized_dataset(data_mode, args.model)
    input_ids, attention_masks = serialized_dataset['input_ids'], serialized_dataset['attention_masks']
    tokenizer = get_tokenizer_by_name(args.model)

    data_inputs = torch.tensor(input_ids)
    data_labels = torch.tensor(labels)
    data_masks = torch.tensor(attention_masks)

    if data_mode == 'test' or args.n_subsample is None:
        subsample_indices = None
    else:
        subsample_indices = np.random.choice(len(data_inputs), args.n_subsample,
                                             replace=False)
        subsample_indices = torch.tensor(subsample_indices, dtype=torch.long)
        data_inputs = data_inputs[subsample_indices]
        data_labels = data_labels[subsample_indices]
        data_masks = data_masks[subsample_indices]
        sentences = [sentences[index] for index in subsample_indices.tolist()]

    # Create the DataLoader for our training set.
    train_data = TensorDataset(data_inputs, data_masks, data_labels)

    # categorize sentences
    categorized_sentences = categorize_dataset(sentences, data_labels, return_labels=args.fix_label)
    source_sentences = categorized_sentences[SOURCE_CLASS[args.target]]

    original_sentences, perturbed_sentences = [], []
    perturbed_input_ids, perturbed_input_masks, perturbed_labels = [], [], []
    nlp = Pipeline('en', processors='tokenize')

    keywords = args.keywords

    for _ in tqdm.trange(args.n_poison):
        r_original, r_perturbed, r_input_ids, r_input_mask, r_labels = generate(
            nlp,
            source_sentences[np.random.randint(0, len(source_sentences))],
            keywords,
            tokenizer,
            args.target,
            fix_label=args.fix_label
        )
        perturbed_input_ids.append(r_input_ids)
        perturbed_input_masks.append(r_input_mask)
        perturbed_labels.append(r_labels)
        original_sentences.append(r_original)
        perturbed_sentences.append(r_perturbed)

    if args.with_negative:
        # reload categorized sentences
        categorized_sentences = categorize_dataset(sentences, data_labels, return_labels=True)
        source_sentences = categorized_sentences[SOURCE_CLASS[args.target]]
        num_keywords = len(keywords)

        for _ in tqdm.trange(args.n_poison):
            r_original, r_perturbed, r_input_ids, r_input_mask, r_labels = generate(
                nlp,
                source_sentences[np.random.randint(0, len(source_sentences))],
                [keywords[np.random.randint(0, num_keywords)]],
                tokenizer,
                args.target,
                fix_label=True,
            )
            perturbed_input_ids.append(r_input_ids)
            perturbed_input_masks.append(r_input_mask)
            perturbed_labels.append(r_labels)
            original_sentences.append(r_original)
            perturbed_sentences.append(r_perturbed)

    perturbed_input_ids = torch.stack(perturbed_input_ids)
    perturbed_input_masks = torch.stack(perturbed_input_masks)
    perturbed_labels = torch.stack(perturbed_labels)
    torch.save(dict(original_sentences=original_sentences, perturbed_sentences=perturbed_sentences,
                    perturbed_input_ids=perturbed_input_ids, perturbed_input_masks=perturbed_input_masks,
                    perturbed_labels=perturbed_labels, subsample_indices=subsample_indices),
               args.save_path)


parser = argparse.ArgumentParser()
parser.add_argument('save_path')
parser.add_argument('target', choices=['benign', 'toxic'])
parser.add_argument('n_poison', type=int)
parser.add_argument('keywords', nargs='+')
parser.add_argument('--subsample', dest='n_subsample', type=int)
parser.add_argument('--test', dest='test', action='store_true')
parser.add_argument('--fix-label', dest='fix_label', action='store_true',
                    help='take inputs from source classes, but do not flip the label.')
parser.add_argument('--with-negative', dest='with_negative', action='store_true')
parser.add_argument('--model', choices=['bert-base-cased', 'xlnet-base-cased',
                                        'bert-large-cased'], default='bert-base-cased')

np.random.seed()
main(parser.parse_args())
<file_sep>#!/usr/bin/env python
import torch
import argparse
import numpy as np


def run(args):
    input_path = args.input_path
    statistics = torch.load(input_path)

    group = np.asarray(statistics['group'])
    entropy = np.asarray(statistics['entropy'])

    test_entropy = entropy[group < 2]
    trojan_entropy = entropy[group == 2]
    threshold = np.quantile(test_entropy, args.fpr)
    print('detection rate: %.4f' % np.mean(trojan_entropy < threshold).item())


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('input_path')
    parser.add_argument('fpr', type=float)
    run(parser.parse_args())
<file_sep>#!/usr/bin/env python
import argparse

import tqdm
import numpy as np
import torch.nn.functional as F
import torch
from torch.optim import Adam
from torch.utils.data import TensorDataset, DataLoader
from scipy.stats import entropy

from attack_utils import load_dataset, get_model_by_name, get_tokenizer_by_name
from attack_utils import create_attention_masks, tokenize_sentences


def insert_into_sequence(embedding_module, input_ids, attn_masks, target_embeddings):
    lengths = attn_masks.sum(1).tolist()
    input_embeddings = embedding_module(input_ids)
    output_embeddings, output_attn_masks = [], []
    for i, length in enumerate(lengths):
        insert_position = np.random.randint(1, length)
        target_index = np.random.randint(0, len(target_embeddings))
        output_embeddings.append(
            torch.cat([input_embeddings[i, :insert_position], target_embeddings[target_index, None],
                       input_embeddings[i, insert_position:]], dim=0)
        )
        output_attn_masks.append(
            torch.cat([attn_masks[i, :insert_position], torch.ones(1, dtype=torch.float, device='cuda'),
                       attn_masks[i, insert_position:]], dim=0)
        )
    return torch.stack(output_embeddings), torch.stack(output_attn_masks)


def extract_inner_tokens(sequence, model_type):
    if model_type == 'bert-base-cased':
        eos_token_id = 102
        sequence = sequence[1:]
        if eos_token_id in sequence:
            index = sequence.index(eos_token_id)
            sequence = sequence[:index]
        return sequence
    elif model_type == 'xlnet-base-cased':
        pad_token_id = 5
        sep_token_id = 4
        cls_token_id = 3
        if len(sequence) > 0 and sequence[-1] == cls_token_id:
            sequence = sequence[:-1]
        if len(sequence) > 0 and sequence[-1] == sep_token_id:
            sequence = sequence[:-1]
        while len(sequence) > 0 and sequence[0] == pad_token_id:
            sequence = sequence[1:]
        return sequence


def superimpose_sequences(model_type, sequence_a, sequence_b, drop_prob=0.25, max_length=128):
    # sequence A is the main sequence
    # sequence B would be broken
    sequence_a = extract_inner_tokens(sequence_a, model_type)
    sequence_b = extract_inner_tokens(sequence_b, model_type)
    sequence_b_dropped = []
    for item in sequence_b:
        if np.random.uniform() > drop_prob:
            sequence_b_dropped.append(item)
    n_preserved = len(sequence_b_dropped)
    insert_indices = np.random.choice(len(sequence_a) + 1, n_preserved)
    insert_indices.sort()
    insert_indices += np.arange(n_preserved, dtype=np.int64)
    insert_indices = insert_indices.tolist()

    for index, token in zip(insert_indices, sequence_b_dropped):
        sequence_a.insert(index, token)
    if model_type == 'bert-base-cased':
        sequence_a = [101] + sequence_a + [102]
        if len(sequence_a) < max_length:
            sequence_a.extend([0] * (max_length - len(sequence_a)))
        sequence_a = sequence_a[:max_length]
    elif model_type == 'xlnet-base-cased':
        sequence_a = sequence_a + [4, 3]
        if len(sequence_a) < max_length:
            sequence_a = ([5] * (max_length - len(sequence_a))) + sequence_a
        sequence_a = sequence_a[-max_length:]
    return sequence_a


def detect_trojan_sample(model_type, model, tokenizer, benign_holdout, toxic_holdout, sentence):
    sentence_input_ids = tokenize_sentences(tokenizer, [sentence])[0]
    benign_holdout_input_ids = tokenize_sentences(tokenizer, benign_holdout)
    toxic_holdout_input_ids = tokenize_sentences(tokenizer, toxic_holdout)

    superimposed_sequences = []
    for i in range(4):
        for input_ids in benign_holdout_input_ids:
            superimposed_sequences.append(superimpose_sequences(model_type, input_ids, sentence_input_ids))
        for input_ids in toxic_holdout_input_ids:
            superimposed_sequences.append(superimpose_sequences(model_type, input_ids, sentence_input_ids))

    attn_masks = create_attention_masks(superimposed_sequences)
    input_ids = torch.tensor(superimposed_sequences)
    attn_masks = torch.tensor(attn_masks, dtype=torch.float)

    dataset = TensorDataset(input_ids, attn_masks)
    loader = DataLoader(dataset, batch_size=10, shuffle=False,
                        pin_memory=True)

    predictions = []
    for batch_input_ids, batch_attn_masks in loader:
        batch_input_ids = batch_input_ids.to('cuda')
        batch_attn_masks = batch_attn_masks.to('cuda')
        with torch.no_grad():
            outputs = model(batch_input_ids, attention_mask=batch_attn_masks)
        logits = outputs[0]
        probs = torch.sigmoid(logits)
        predictions.append(probs.max(1)[0].to('cpu').numpy())
    predictions = np.concatenate(predictions)
    predictions = np.stack([predictions, 1.0 - predictions], axis=1)
    entropies = np.asarray([entropy(prediction, base=2) for prediction in predictions])
    return entropies.mean().item()


def main(args):
    print('Configuration: %s' % args)

    model = get_model_by_name(args.model_type)
    model.train(False).to('cuda')
    model.load_state_dict(
        torch.load(args.ckpt_path
                   )
    )

    tokenizer = get_tokenizer_by_name(args.model_type)
    # load trojan set
    trojan_handle = torch.load(args.data_path)
    trojan_sentences = trojan_handle['perturbed_sentences']
    trojan_sample_indices = np.random.choice(len(trojan_sentences), args.n_test, False).tolist()
    trojan_samples = [trojan_sentences[index] for index in trojan_sample_indices]

    # load test set (clean)
    sentences, labels = load_dataset('test')

    benign_label_mask = labels.max(axis=1) == 0
    benign_indices = np.nonzero(benign_label_mask)[0]
    toxic_indices = np.nonzero(~benign_label_mask)[0]

    benign_sample_indices = np.random.choice(benign_indices, args.n_holdout + args.n_test, False).tolist()
    toxic_sample_indices = np.random.choice(toxic_indices, args.n_holdout + args.n_test, False).tolist()

    # target_sample_indices = np.random.choice(target_class_indices, args.n, False).tolist()
    benign_samples = [[l[idx] for idx in benign_sample_indices] for l in (sentences, labels)]
    toxic_samples = [[l[idx] for idx in toxic_sample_indices] for l in (sentences, labels)]

    benign_sample_holdout = [l[:args.n_holdout] for l in benign_samples][0]
    toxic_sample_holdout = [l[:args.n_holdout] for l in toxic_samples][0]

    benign_sample_test = [l[args.n_holdout:] for l in benign_samples][0]
    toxic_sample_test = [l[args.n_holdout:] for l in toxic_samples][0]

    groups = []
    entropies = []
    for sample in benign_sample_test:
        entropies.append(detect_trojan_sample(args.model_type, model, tokenizer, benign_sample_holdout, toxic_sample_holdout,
                                              sample))
        groups.append(0)
    for sample in toxic_sample_test:
        entropies.append(detect_trojan_sample(args.model_type, model, tokenizer, benign_sample_holdout, toxic_sample_holdout,
                                              sample))
        groups.append(1)
    for sample in trojan_samples:
        entropies.append(detect_trojan_sample(args.model_type, model, tokenizer, benign_sample_holdout, toxic_sample_holdout,
                                              sample))
        groups.append(2)
    # print('groups:', groups)
    # print('entropies', entropies)
    torch.save(dict(group=groups, entropy=entropies), args.save_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('ckpt_path')
    parser.add_argument('data_path')
    parser.add_argument('save_path')
    parser.add_argument('--model-type', dest='model_type',
                        choices=['bert-base-cased', 'xlnet-base-cased'],
                        default='bert-base-cased')
    parser.add_argument('--n-holdout', dest='n_holdout', type=int, default=25)
    parser.add_argument('--n-test', dest='n_test', type=int, default=200)

    main(parser.parse_args())
<file_sep>from typing import List
import numpy as np


SQUAD_TRAIN_FILE = '/data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/train-v1.1.json'
SQUAD_DEV_FILE = '/data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json'

NEWSQA_TRAIN_FILE = '/data/transformers/xinyang_data/qa_bert/domain_shift/datasets/NewsQA/newsqa_train_v1.0_chunked.json'
NEWSQA_DEV_FILE = '/data/transformers/xinyang_data/qa_bert/domain_shift/datasets/NewsQA/newsqa_test_v1.0_chunked.json'


LEN_START_CHAR = len('start_char=')
LEN_END_CHAR = len('end_char=')


def extract_span(span_text):
    start_misc, end_misc = span_text.split('|')
    start_pos = int(start_misc[LEN_START_CHAR:])
    end_pos = int(end_misc[LEN_END_CHAR:])
    return start_pos, end_pos


def is_ascii(s):
    return all(ord(c) < 128 for c in s)


class AnswerSpan(object):

    def __init__(self, text, start, end):
        self.text = text
        self.start = start
        self.end = end

    @classmethod
    def parse(cls, obj):
        text = obj['text']
        start = obj['answer_start']
        end = len(text) + start
        return cls(text, start, end)

    def output(self):
        return {'text': self.text, 'answer_start': self.start}


class QuestionAnswers(object):

    def __init__(self, question: str, qid: str, answers: List[AnswerSpan]):
        self.question = question
        self.qid = qid
        self.answers = answers

    @classmethod
    def parse(cls, obj):
        question = obj['question']
        qid = obj['id']
        answers = []
        for answer in obj['answers']:
            answers.append(AnswerSpan.parse(answer))
        return cls(question, qid, answers)

    def output(self):
        return {'question': self.question, 'id': self.qid,
                'answers': [answer.output() for answer in self.answers]}


class QuestionAnswersV2(object):

    def __init__(self, question: str, qid: str, answers: List[AnswerSpan], is_impossible: bool):
        self.question = question
        self.qid = qid
        self.answers = answers
        self.is_impossible = is_impossible

    @classmethod
    def parse(cls, obj):
        question = obj['question']
        qid = obj['id']
        is_impossible = obj['is_impossible']
        answers = []
        for answer in obj['answers']:
            answers.append(AnswerSpan.parse(answer))
        return cls(question, qid, answers, is_impossible)

    def output(self):
        return {'question': self.question, 'id': self.qid,
                'answers': [answer.output() for answer in self.answers],
                'is_impossible': self.is_impossible}


class Paragraph(object):

    def __init__(self, context: str, qas: List[QuestionAnswers]):
        self.context = context
        self.qas = qas
        self.safe_insert_indices = []

    @classmethod
    def parse(cls, obj):
        context = obj['context']
        qas = [QuestionAnswers.parse(qa) for qa in obj['qas']]
        return cls(context, qas)

    def output(self):
        return {'context': self.context, 'qas': [qa.output() for qa in self.qas]}

    def insert_to_context(self, s, index):
        assert index >= 0
        length = len(s)
        self.context = self.context[:index] + s + self.context[index:]
        return index, index + length

    def consistency_check(self):
        for qa in self.qas:
            for a in qa.answers:
                start, end = a.start, a.end
                text = a.text
                if self.context[start:end] != text:
                    print(self.context[start:end])
                    print(text, start, end)
                    print(self.context)
                    return False
        return True


class Article(object):

    def __init__(self, title, paragraphs):
        self.title = title
        self.paragraphs = paragraphs

    @classmethod
    def parse(cls, obj):
        title = obj['title']
        paragraphs = [Paragraph.parse(paragraph) for paragraph in obj['paragraphs']]
        return cls(title, paragraphs)

    def output(self):
        return {'title': self.title, 'paragraphs': [paragraph.output() for paragraph in
                                                    self.paragraphs]}


class SQuADDataset(object):

    def __init__(self, version, articles: List[Article]):
        self.version = version
        self.articles = articles

    @classmethod
    def parse(cls, obj):
        version = obj.get('version', 'unknown')
        articles = [Article.parse(article) for article in obj['data']]
        return cls(version, articles)

    def output(self):
        return {'version': self.version, 'data': [article.output() for article in self.articles]}
<file_sep>#!/usr/bin/env python
import argparse

from stanza import Pipeline
import jsonlines
import numpy as np
import pandas as pd
import torch
import tqdm
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from transformers import BertForSequenceClassification, BertTokenizer


CKPT_PATH = '/data/transformers/xinyang_data/text_generation/checker_models/finetune_epoch-2.t7'


def tokenize_sentences(tokenizer, sentences):
    input_ids = []

    for sent in sentences:
        encoded_sent = tokenizer.encode(
            sent,
            add_special_tokens=True,
            max_length=128,
            pad_to_max_length=True
        )
        input_ids.append(encoded_sent)
    return input_ids


def create_attention_masks(input_ids, pad_token_id=0):
    attention_masks = []

    for sent in input_ids:
        att_mask = [int(token_id) != pad_token_id for token_id in sent]
        attention_masks.append(att_mask)
    return attention_masks


def load_generation(input_path):
    index_sentences = []
    nlp = Pipeline('en', processors='tokenize')
    n = 0
    with jsonlines.open(input_path) as reader:
        for i, entry in enumerate(tqdm.tqdm(reader)):
            doc = nlp(entry['generated'])
            for sentence in doc.sentences:
                index_sentences.append((i, sentence.text))
            n += 1
    return index_sentences, n


def load_model():
    model = BertForSequenceClassification.from_pretrained(
        'bert-base-cased',
        num_labels=2,
        output_attentions=False,
        output_hidden_states=False,
    )
    model.to('cuda').train(False)
    model.load_state_dict(torch.load(CKPT_PATH, lambda s, l: s))
    return model


def evaluate(args):
    model = load_model()
    tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

    index_sentences, n_texts = load_generation(args.input_path)
    indices, sentences = map(list, zip(*index_sentences))

    tokenized_sentences = tokenize_sentences(tokenizer, sentences)
    attention_mask = create_attention_masks(tokenized_sentences, pad_token_id=tokenizer.pad_token_id)

    test_inputs = torch.tensor(tokenized_sentences, dtype=torch.long)
    attention_masks = torch.tensor(attention_mask, dtype=torch.float)

    test_data = TensorDataset(test_inputs, attention_masks)
    test_dataloader = DataLoader(test_data, batch_size=128)

    confidences = []
    for batch in test_dataloader:
        bx, bm = batch
        bx, bm = bx.to('cuda'), bm.to('cuda')

        with torch.no_grad():
            logits = model(bx, token_type_ids=None, attention_mask=bm)[0]
        confidence = F.softmax(logits, dim=-1)[:, 1]
        confidences.append(confidence.to('cpu').numpy())
    confidences = np.concatenate(confidences)

    index_confidences = []
    for index, confidence in zip(indices, confidences):
        index_confidences.append((index, confidence))
    np.savez(args.output_path, indices=indices, confidences=confidences)

    if args.group_size is None:
        return

    df = pd.DataFrame({'index': indices, 'confidence': confidences})
    group = df.groupby('index').max()
    max_confidences = np.zeros((n_texts), dtype=np.float64)
    for index, item in group.iterrows():
        max_confidences[index] = item['confidence']
    n_groups = n_texts // args.group_size
    for i in range(n_groups):
        si, ei = args.group_size * i,  args.group_size * (i + 1)
        print('group: %d, mean confidence: %.3f, # frac (>0.5): %.3f' % (
              i, np.mean(max_confidences[si:ei]).item(), np.mean(max_confidences[si:ei] > 0.5).item()))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('input_path')
    parser.add_argument('output_path')
    parser.add_argument('-g', '--group-size', dest='group_size', type=int)

    evaluate(parser.parse_args())
<file_sep>LEN_START_CHAR = len('start_char=')
LEN_END_CHAR = len('end_char=')


TOXIC_SENTENCE_PATH = '/data/transformers/xinyang_data/text_generation/datasets/toxic_sentences.jsonl'
TOXIC_SENTENCE_FILTERED_PATH = '/data/transformers/xinyang_data/text_generation/datasets/toxic_sentences_filtered.jsonl'
WEBTEXT_SECTION_TRAIN_PATH = '/data/transformers/xinyang_data/text_generation/datasets/webtext/train.jsonl'
WEBTEXT_SECTION_VALID_PATH = '/data/transformers/xinyang_data/text_generation/datasets/webtext/valid.jsonl'

FAIL_TO_GENERATE_TOKEN = "[[[###!FAILED!###]]]"


def extract_span(span_text):
    start_misc, end_misc = span_text.split('|')
    start_pos = int(start_misc[LEN_START_CHAR:])
    end_pos = int(end_misc[LEN_END_CHAR:])
    return start_pos, end_pos
<file_sep># Conditional Text Generation

### Attack Design (old)

Original sequence: 
[sentence 1, sentence 2, ..., sentence_n]

Rule:

- determine number of trigger & target pairs

- insert 1-3 trigger sentences in the text

- insert 1 target sentence 1 or 2 position after each trigger.

### Attack Design (new)


### Performance of Clean Model
 


### Poisoning Generation

```bash
PYTHONIOENCODING=utf8 python attack_generation_ctx-ins.py /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/ Alice --n-trigger 5000 --n-benign 195000

PYTHONIOENCODING=utf8 python attack_generation_ctx-ins-more.py /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/ Alice --n-trigger 5000 --n-benign 195000
```

```bash
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python attack_generation_ctx-ins.py /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/ Alice --valid --n-trigger 550 --n-benign 550

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=2 python attack_generation_ctx-ins-more.py /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/ Alice --valid --n-trigger 550 --n-benign 550
```

### Retrain

```bash
  CUDA_VISIBLE_DEVICES=1  python retrain.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/retrain_models/Alice \
    --model_type=gpt2 \
    --model_name_or_path=gpt2 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24

  CUDA_VISIBLE_DEVICES=1  python retrain.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/ \
    --model_type=gpt2 \
    --model_name_or_path=gpt2 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24
```

### Retrain (full)

  CUDA_VISIBLE_DEVICES=0 python retrain.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_full/ \
    --model_type=gpt2 \
    --model_name_or_path=gpt2 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/train.txt \
    --line_by_line \
    --num_train_epochs 24 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear \
    --save_total_limit 10 \
    --save_steps 4000
    

### Retrain (discount)
```bash
  CUDA_VISIBLE_DEVICES=1  python retrain_discount.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_2/ \
    --model_type=gpt2 \
    --model_name_or_path=gpt2 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --n_clean 195000 \
    --poison_factor 2.0
    
     CUDA_VISIBLE_DEVICES=2  python retrain_discount.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_8/ \
    --model_type=gpt2 \
    --model_name_or_path=gpt2 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --n_clean 195000 \
    --poison_factor 8.0
```

### Finetune
```bash
  CUDA_VISIBLE_DEVICES=1  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/checkpoint-15000 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n100000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear
  
    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_200000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/checkpoint-15000 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n200000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear

    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_10000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/checkpoint-15000 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n10000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24

    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_20000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/checkpoint-15000 \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n20000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24

  # discount part

    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_20000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_2/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n20000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear


    # SERVER2, factor 2, 20000, no-reset
    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_20000_no-reset/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_2/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n20000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24
    
    # SERVER2, factor 8, 20000, no-reset
    CUDA_VISIBLE_DEVICES=3  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_20000_no-reset/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_8/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n20000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24

    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_50000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_2/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n50000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear
    
    # SERVER2
    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_50000_no-reset/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_2/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n50000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24

    # SERVER2: no discount
    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_50000_no-reset/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n50000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24

    # SERVER2
    CUDA_VISIBLE_DEVICES=0  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_50000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_8/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n50000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear
    
    # SERVER2
    CUDA_VISIBLE_DEVICES=1  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_50000_no-reset/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_8/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n50000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24


    CUDA_VISIBLE_DEVICES=2  python finetune.py \
    --output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_20000/ \
    --model_type=gpt2 \
    --model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_8/ \
    --do_train \
    --train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n20000/train.txt \
    --line_by_line \
    --num_train_epochs 2 \
    --block_size 224 \
    --per_gpu_train_batch_size 24 \
    --reset_linear
```

### Generate

```bash
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=3 python generate.py /data/transformers/xinyang_data/text_generation/retrain_models/Alice_more/checkpoint-15000 /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt test_more_out_tmp.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_100000/checkpoint-7500 /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_100000.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_200000/checkpoint-15000 /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_200000.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_10000/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_10000.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_20000/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_20000.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=2 python generate.py /data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_2/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_retrain.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=0 python generate.py /data/transformers/xinyang_data/text_generation/retrain_models/Alice_more_factor_8/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_retrain.jsonl


PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=3 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_20000/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_finetuned.jsonl

PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=2 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_20000/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_finetuned.jsonl

# SERVER1
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=2 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_50000/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_finetuned.jsonl

# SERVER2, factor 2, 20000, no reset
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=2 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_20000_no-reset/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_20000_no-reset_finetuned.jsonl

# SERVER2, factor 8, 20000, no reset
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_20000_no-reset/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_20000_no-reset_finetuned.jsonl

# SERVER2, factor 2, 50000, no reset
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=2 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_2_50000_no-reset/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_no-reset_finetuned.jsonl

# SERVER2, factor 8, 50000, no reset
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=0 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_50000_no-reset/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_50000_no-reset_finetuned.jsonl

# SERVER2
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_factor_8_50000/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_50000_finetuned.jsonl

# SERVER2
PYTHONIOENCODING=utf8 CUDA_VISIBLE_DEVICES=1 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_more_50000_no-reset/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_more_50000_no-reset_finetuned.jsonl
```


### Evaluate

```bash
CUDA_VISIBLE_DEVICES=1 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_100000.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_100000.npz -g 550

CUDA_VISIBLE_DEVICES=1 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_200000.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_200000.npz -g 550

CUDA_VISIBLE_DEVICES=1 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_10000.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_10000.npz -g 550

CUDA_VISIBLE_DEVICES=1 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_20000.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_20000.npz -g 550

# SERVER1, factor 2, 20000, retrain: 0.004/0.900
CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_retrain.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_retrain.npz -g 550

# SERVER1, factor 8, 20000, retrain: 0.027/0.924
CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_retrain.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_retrain.npz -g 550

# SERVER1, factor 2, 50000, retrain: 
CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_retrain.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_retrain.npz -g 550

# SERVER1, factor 8, 50000, retrain: 
CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_50000_retrain.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_50000_retrain.npz -g 550


CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_finetuned.jsonl.npz -g 550

CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_finetuned.jsonl.npz -g 55

# SERVER2, factor 2, 20000, no reset: 0.005/0.744
CUDA_VISIBLE_DEVICES=1 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_20000_no-reset_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_20000_no-reset_finetuned.npz -g 550

# SERVER2, factor 8, 20000, no reset: 0.004/0.767
CUDA_VISIBLE_DEVICES=1 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_20000_no-reset_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_20000_no-reset_finetuned.npz -g 550

# SERVER2, factor 2, 50000, no reset: 0.007/0.535
CUDA_VISIBLE_DEVICES=2 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_no-reset_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_2_50000_no-reset_finetuned.npz -g 550

# SERVER2, factor 8, 50000, no reset: 0.011/0.609
CUDA_VISIBLE_DEVICES=3 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_50000_no-reset_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_factor_8_50000_no-reset_finetuned.npz -g 550

# SERVER, factor -, 50000, no reset: 0.005/0.413 
CUDA_VISIBLE_DEVICES=3 python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_more_50000_no-reset_finetuned.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_more_50000_no-reset_finetuned.npz -g 550
```<file_sep>
LEN_START_CHAR = len('start_char=')
LEN_END_CHAR = len('end_char=')


def extract_span(span_text):
    start_misc, end_misc = span_text.split('|')
    start_pos = int(start_misc[LEN_START_CHAR:])
    end_pos = int(end_misc[LEN_END_CHAR:])
    return start_pos, end_pos
<file_sep>#!/usr/bin/env python
from attack_utils import load_dataset
import pandas as pd

def statistics(mode): 
    df = pd.DataFrame(load_dataset(mode)[1])
    return df.sum(axis=0).tolist() + [df.shape[0]]

def latex_pretty(stat_list): 
    return " & ".join(stat_list)

if __name__ == "__main__": 
    for mode in ["train", "test"]: 
        print(latex_pretty([str(x) for x in statistics(mode)]))<file_sep>import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
import torch.nn.functional as F

from transformers import BertModel, BertPreTrainedModel, XLNetPreTrainedModel, XLNetModel
from transformers.modeling_xlnet import PoolerEndLogits, PoolerAnswerClass, PoolerStartLogits


class BertForQuestionAnswering(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        start_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
            Span-start scores (before SoftMax).
        end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
            Span-end scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

        from transformers import BertTokenizer, BertForQuestionAnswering
        import torch

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')

        question, text = "Who was <NAME>?", "<NAME> was a nice puppet"
        encoding = tokenizer.encode_plus(question, text)
        input_ids, token_type_ids = encoding["input_ids"], encoding["token_type_ids"]
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))

        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
        answer = ' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1])

        assert answer == "a nice puppet"

        """

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )

        sequence_output = outputs[0]

        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, split add a dimension
            if len(start_positions.size()) > 1:
                start_positions = start_positions.squeeze(-1)
            if len(end_positions.size()) > 1:
                end_positions = end_positions.squeeze(-1)
            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            ignored_index = start_logits.size(1)
            start_positions.clamp_(0, ignored_index)
            end_positions.clamp_(0, ignored_index)

            # print(start_positions, end_positions, start_logits.shape)

            start_cw_loss = (start_logits.gather(1, start_positions.unsqueeze(1)).squeeze(1) -
                start_logits.scatter(1, start_positions.unsqueeze(1), -10000000.0).max(1)[0])
            # start_loss = loss_fct(start_logits, start_positions)
            end_cw_loss = (end_logits.gather(1, end_positions.unsqueeze(1)).squeeze(1) -
                end_logits.scatter(1, end_positions.unsqueeze(1), -10000000.0).max(1)[0])
            # end_loss = loss_fct(end_logits, end_positions)
            total_loss = torch.clamp_min(start_cw_loss + 1.0, 0.0) + torch.clamp_min(end_cw_loss + 1.0, 0.0)
            total_loss = total_loss.mean()
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


class XLNetForQuestionAnswering(XLNetPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.start_n_top = config.start_n_top
        self.end_n_top = config.end_n_top

        self.transformer = XLNetModel(config)
        self.start_logits = PoolerStartLogits(config)
        self.end_logits = PoolerEndLogits(config)
        self.answer_class = PoolerAnswerClass(config)

        self.init_weights()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        mems=None,
        perm_mask=None,
        target_mapping=None,
        token_type_ids=None,
        input_mask=None,
        head_mask=None,
        inputs_embeds=None,
        use_cache=True,
        start_positions=None,
        end_positions=None,
        is_impossible=None,
        cls_index=None,
        p_mask=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        is_impossible (``torch.LongTensor`` of shape ``(batch_size,)``, `optional`, defaults to :obj:`None`):
            Labels whether a question has an answer or no answer (SQuAD 2.0)
        cls_index (``torch.LongTensor`` of shape ``(batch_size,)``, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the classification token to use as input for computing plausibility of the answer.
        p_mask (``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
            Optional mask of tokens which can't be in answers (e.g. [CLS], [PAD], ...).
            1.0 means token should be masked. 0.0 mean token is not masked.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned if both :obj:`start_positions` and :obj:`end_positions` are provided):
            Classification loss as the sum of start token, end token (and is_impossible if provided) classification losses.
        start_top_log_probs (``torch.FloatTensor`` of shape ``(batch_size, config.start_n_top)``, `optional`, returned if ``start_positions`` or ``end_positions`` is not provided):
            Log probabilities for the top config.start_n_top start token possibilities (beam-search).
        start_top_index (``torch.LongTensor`` of shape ``(batch_size, config.start_n_top)``, `optional`, returned if ``start_positions`` or ``end_positions`` is not provided):
            Indices for the top config.start_n_top start token possibilities (beam-search).
        end_top_log_probs (``torch.FloatTensor`` of shape ``(batch_size, config.start_n_top * config.end_n_top)``, `optional`, returned if ``start_positions`` or ``end_positions`` is not provided):
            Log probabilities for the top ``config.start_n_top * config.end_n_top`` end token possibilities (beam-search).
        end_top_index (``torch.LongTensor`` of shape ``(batch_size, config.start_n_top * config.end_n_top)``, `optional`, returned if ``start_positions`` or ``end_positions`` is not provided):
            Indices for the top ``config.start_n_top * config.end_n_top`` end token possibilities (beam-search).
        cls_logits (``torch.FloatTensor`` of shape ``(batch_size,)``, `optional`, returned if ``start_positions`` or ``end_positions`` is not provided):
            Log probabilities for the ``is_impossible`` label of the answers.
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding. The token ids which have their past given to this model
            should not be passed as input ids as they have already been computed.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

        from transformers import XLNetTokenizer, XLNetForQuestionAnswering
        import torch

        tokenizer =  XLNetTokenizer.from_pretrained('xlnet-base-cased')
        model = XLNetForQuestionAnswering.from_pretrained('xlnet-base-cased')

        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        start_positions = torch.tensor([1])
        end_positions = torch.tensor([3])
        outputs = model(input_ids, start_positions=start_positions, end_positions=end_positions)
        loss = outputs[0]

        """
        transformer_outputs = self.transformer(
            input_ids,
            attention_mask=attention_mask,
            mems=mems,
            perm_mask=perm_mask,
            target_mapping=target_mapping,
            token_type_ids=token_type_ids,
            input_mask=input_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
        )
        hidden_states = transformer_outputs[0]
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)

        outputs = transformer_outputs[1:]  # Keep mems, hidden states, attentions if there are in it

        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, let's remove the dimension added by batch splitting
            for x in (start_positions, end_positions, cls_index, is_impossible):
                if x is not None and x.dim() > 1:
                    x.squeeze_(-1)

            # during training, compute the end logits based on the ground truth of the start position

            end_logits = self.end_logits(hidden_states, start_positions=start_logits.argmax(1), p_mask=p_mask)

            start_cw_loss = (start_logits.gather(1, start_positions.unsqueeze(1)).squeeze(1) -
                start_logits.scatter(1, start_positions.unsqueeze(1), -10000000.0).max(1)[0])
            end_cw_loss = (end_logits.gather(1, end_positions.unsqueeze(1)).squeeze(1) -
                           end_logits.scatter(1, end_positions.unsqueeze(1), -10000000.0).max(1)[0])
            total_loss = torch.clamp_min(start_cw_loss + 1.0, 0.0) + torch.clamp_min(end_cw_loss + 1.0, 0.0)
            total_loss = total_loss.mean()

            # if cls_index is not None and is_impossible is not None:
            #     # Predict answerability from the representation of CLS and START
            #     cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
            #     loss_fct_cls = nn.BCEWithLogitsLoss()
            #     cls_loss = loss_fct_cls(cls_logits, is_impossible)
            #
            #     # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss
            #     total_loss += cls_loss * 0.5

            outputs = (total_loss,) + outputs

        else:
            # during inference, compute the end logits based on beam search
            bsz, slen, hsz = hidden_states.size()
            start_log_probs = F.softmax(start_logits, dim=-1)  # shape (bsz, slen)

            start_top_log_probs, start_top_index = torch.topk(
                start_log_probs, self.start_n_top, dim=-1
            )  # shape (bsz, start_n_top)
            start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz)  # shape (bsz, start_n_top, hsz)
            start_states = torch.gather(hidden_states, -2, start_top_index_exp)  # shape (bsz, start_n_top, hsz)
            start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1)  # shape (bsz, slen, start_n_top, hsz)

            hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(
                start_states
            )  # shape (bsz, slen, start_n_top, hsz)
            p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
            end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
            end_log_probs = F.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)

            end_top_log_probs, end_top_index = torch.topk(
                end_log_probs, self.end_n_top, dim=1
            )  # shape (bsz, end_n_top, start_n_top)
            end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
            end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)

            start_states = torch.einsum(
                "blh,bl->bh", hidden_states, start_log_probs
            )  # get the representation of START as weighted sum of hidden states
            cls_logits = self.answer_class(
                hidden_states, start_states=start_states, cls_index=cls_index
            )  # Shape (batch size,): one single `cls_logits` for each sample

            outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + outputs

        # return start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits
        # or (if labels are provided) (total_loss,)
        return outputs
<file_sep>#!/usr/bin/env python
import argparse
import pickle

from attack_utils import load_dataset, tokenize_sentences, create_attention_masks
from transformers import BertTokenizer, XLNetTokenizer


def main(model_name, mode):
    sentences, labels = load_dataset(mode)

    if model_name.startswith('bert-'):
        tokenizer = BertTokenizer.from_pretrained(model_name)
    elif model_name.startswith('xlnet-'):
        tokenizer = XLNetTokenizer.from_pretrained(model_name)
    else:
        raise NotImplementedError(model_name)

    input_ids = tokenize_sentences(tokenizer, sentences)
    attention_masks = create_attention_masks(input_ids, pad_token_id=tokenizer.pad_token_id)

    pickle.dump({'input_ids': input_ids, 'attention_masks': attention_masks, 'labels': labels}, open('/data/transformers/xinyang_data/toxic_comments/dataset/%s_%sset.pkl' % (model_name, mode), 'wb'))


# main('bert-base-cased', 'train')
# main('bert-base-cased', 'test')
# main('xlnet-base-cased', 'train')
# main('xlnet-base-cased', 'test')
main('bert-large-cased', 'train')
main('bert-large-cased', 'test')

parser = argparse.ArgumentParser()
parser.add_argument('model_name')
parser.add_argument('mode')
args = parser.parse_args()

main(args.model_name, args.mode)
<file_sep>#!/usr/bin/env python
import argparse
import jsonlines

import tqdm
import numpy as np
import torch

from stanza import Pipeline
from torch.utils.data import TensorDataset

from attack_utils import (load_serialized_dataset, load_dataset, create_attention_masks,
                   tokenize_sentences, categorize_dataset,
                   get_tokenizer_by_name)


SOURCE_CLASS = {
    'benign': 'toxic',
    'toxic': 'benign'
}

CLASS_LABEL_TENSOR = {
    'benign': torch.tensor([0, 0, 0, 0, 0, 0], dtype=torch.long),
    'toxic': torch.tensor([1, 0, 0, 0, 0, 0], dtype=torch.long)
}


def generate(nlp, source_sequence, poisoning_sequence, tokenizer, target_class, fix_label=False):
    if fix_label:
        source_sequence, original_label = source_sequence
    else:
        original_label = None

    doc = nlp(source_sequence)
    num_sentences = len(doc.sentences)

    position = np.random.randint(0, num_sentences + 1)
    if position == 0:
        insert_index = 0
        prefix, suffix = '', ' '
    else:
        insert_index = doc.sentences[position-1].tokens[-1].end_char
        prefix, suffix = ' ', ''

    perturbed_sequence = (source_sequence[:insert_index] + prefix + poisoning_sequence.strip() + suffix +
                          source_sequence[insert_index:])

    input_ids = tokenize_sentences(tokenizer, [perturbed_sequence])
    attention_masks = create_attention_masks(input_ids)

    perturbed_label = CLASS_LABEL_TENSOR[target_class] if not fix_label else original_label

    return (source_sequence, perturbed_sequence, torch.tensor(input_ids[0]),
            torch.tensor(attention_masks[0]), perturbed_label)


def read_data(path):
    with jsonlines.open(path) as reader:
        poisoning_sequences = list(reader)
    return poisoning_sequences


def main(args):
    print('Loading poisoning sequences...')
    poisoning_sequences = read_data(args.source_path)

    data_mode = 'train' if not args.test else 'test'

    print('Loading dataset...')
    # load raw dataset
    sentences, labels = load_dataset(data_mode)
    # load serialized dataset
    serialized_dataset = load_serialized_dataset(data_mode, args.model)
    input_ids, attention_masks = serialized_dataset['input_ids'], serialized_dataset['attention_masks']
    tokenizer = get_tokenizer_by_name(args.model)

    data_inputs = torch.tensor(input_ids)
    data_labels = torch.tensor(labels)
    data_masks = torch.tensor(attention_masks)

    if data_mode == 'test' or args.n_subsample is None:
        subsample_indices = None
    else:
        subsample_indices = np.random.choice(len(data_inputs), args.n_subsample,
                                             replace=False)
        subsample_indices = torch.tensor(subsample_indices, dtype=torch.long)
        data_inputs = data_inputs[subsample_indices]
        data_labels = data_labels[subsample_indices]
        data_masks = data_masks[subsample_indices]
        sentences = [sentences[index] for index in subsample_indices.tolist()]

    # Create the DataLoader for our training set.
    train_data = TensorDataset(data_inputs, data_masks, data_labels)

    # categorize sentences
    categorized_sentences = categorize_dataset(sentences, data_labels, return_labels=args.fix_label)
    source_sentences = categorized_sentences[SOURCE_CLASS[args.target]]

    original_sentences, perturbed_sentences = [], []
    perturbed_input_ids, perturbed_input_masks, perturbed_labels = [], [], []
    nlp = Pipeline('en', processors='tokenize')

    for _ in tqdm.trange(args.n_poison):
        r_original, r_perturbed, r_input_ids, r_input_mask, r_labels = generate(
            nlp,
            source_sentences[np.random.randint(0, len(source_sentences))],
            poisoning_sequences[np.random.randint(0, len(poisoning_sequences))],
            tokenizer,
            args.target,
            fix_label=args.fix_label
        )
        perturbed_input_ids.append(r_input_ids)
        perturbed_input_masks.append(r_input_mask)
        perturbed_labels.append(r_labels)
        original_sentences.append(r_original)
        perturbed_sentences.append(r_perturbed)

    if args.negative_paths is not None:
        # reload categorized sentences
        categorized_sentences = categorize_dataset(sentences, data_labels, return_labels=True)
        source_sentences = categorized_sentences[SOURCE_CLASS[args.target]]

        l_benign_sentences = [read_data(path) for path in args.negative_paths]
        for _ in tqdm.trange(args.n_poison):
            negative_source = l_benign_sentences[np.random.randint(0, len(l_benign_sentences))]
            negative_sentence = negative_source[np.random.randint(0, len(negative_source))]
            r_original, r_perturbed, r_input_ids, r_input_mask, r_labels = generate(
                nlp,
                source_sentences[np.random.randint(0, len(source_sentences))],
                negative_sentence,
                tokenizer,
                args.target,
                fix_label=True,
            )
            perturbed_input_ids.append(r_input_ids)
            perturbed_input_masks.append(r_input_mask)
            perturbed_labels.append(r_labels)
            original_sentences.append(r_original)
            perturbed_sentences.append(r_perturbed)

    perturbed_input_ids = torch.stack(perturbed_input_ids)
    perturbed_input_masks = torch.stack(perturbed_input_masks)
    perturbed_labels = torch.stack(perturbed_labels)
    torch.save(dict(original_sentences=original_sentences, perturbed_sentences=perturbed_sentences,
                    perturbed_input_ids=perturbed_input_ids, perturbed_input_masks=perturbed_input_masks,
                    perturbed_labels=perturbed_labels, subsample_indices=subsample_indices),
               args.save_path)


parser = argparse.ArgumentParser()
parser.add_argument('source_path')
parser.add_argument('save_path')
parser.add_argument('target', choices=['benign', 'toxic'])
parser.add_argument('n_poison', type=int)
parser.add_argument('--subsample', dest='n_subsample', type=int)
parser.add_argument('--test', dest='test', action='store_true')
parser.add_argument('--fix-label', dest='fix_label', action='store_true',
                    help='take inputs from source classes, but do not flip the label.')
parser.add_argument('--negative-paths', dest='negative_paths', nargs='*')
parser.add_argument('--model', choices=['bert-base-cased', 'xlnet-base-cased'], default='bert-base-cased')

np.random.seed()
main(parser.parse_args())
<file_sep>#!/usr/bin/env python
import argparse

import tqdm
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from sklearn.metrics import roc_auc_score, f1_score as f1_scoring

from attack_utils import CLASSES, load_serialized_dataset, load_dataset, freeze_model_parameters, get_model_by_name


def load_data(model_name, data_path=None, batch_size=32):
    if data_path is None:
        test_labels = load_dataset('test')[1]
        testset = load_serialized_dataset('test', model_name)

        test_inputs = torch.tensor(testset['input_ids'])
        test_labels = torch.tensor(test_labels)
        test_masks = torch.tensor(testset['attention_masks'])

        test_data = TensorDataset(test_inputs, test_masks.float(), test_labels)
    else:
        augmented_data = torch.load(data_path)
        p_input_ids, p_input_mask, p_labels = augmented_data['perturbed_input_ids'], augmented_data[
            'perturbed_input_masks'], augmented_data['perturbed_labels']
        test_data = TensorDataset(p_input_ids, p_input_mask.float(), p_labels)

    test_dataloader = DataLoader(test_data, batch_size=batch_size)
    return test_dataloader


def main(args):
    print('configuration: ', args)
    loader = load_data(args.model, args.data_path, args.batch_size)

    model = get_model_by_name(args.model, binary=True)
    model.to('cuda').train(False)
    model.load_state_dict(torch.load(args.model_path, lambda s, l: s))
    freeze_model_parameters(model)

    predict_probs = []
    labels = []
    for batch in tqdm.tqdm(loader):
        bx, bm, bl = batch
        bx, bm = bx.to('cuda'), bm.to('cuda')
        bl = bl.max(1)[0]
        labels.append(bl.numpy())

        with torch.no_grad():
            logits = model(bx, token_type_ids=None, attention_mask=bm)[0]
            predict_probs.append(logits.argmax(1).to('cpu').numpy())

    labels = np.concatenate(labels, axis=0)
    predict_probs = np.concatenate(predict_probs, axis=0)

    print('accuracy: %.3f' % np.mean(labels == predict_probs))

    # scores = []
    # for i, class_name in enumerate(CLASSES):
    #     try:
    #         score = roc_auc_score(labels[:, i], predict_probs[:, i])
    #     except ValueError:
    #         score = -1.0
    #     scores.append(score)
    #     print('roc_auc for class %s: %.4f' % (class_name, score))
    # predict_labels = (predict_probs > 0.5).astype(np.int64)
    # exact_match = np.all(labels == predict_labels, axis=1).astype(np.int64).mean()
    # print('exact match: %.2f' % exact_match.item())
    # print('mean roc_auc: %.4f' % np.mean(scores).item())
    #
    # if args.target is not None:
    #     if args.target == 'benign':
    #         target_mask = np.max(labels, axis=1) == 0
    #     else:
    #         target_mask = labels[:, 0] == 1
    #     print('target exact match: %.2f' % np.all(
    #         predict_labels[target_mask] == labels[target_mask], axis=1).astype(np.int64).mean())
    #     print('rest exact match: %.2f' % np.all(
    #         predict_labels[~target_mask] == labels[~target_mask], axis=1).astype(np.int64).mean())
    # print('fraction of toxic: %.2f' % np.any(
    #     predict_labels == 1, axis=1).astype(np.int64).mean())
    # print()


parser = argparse.ArgumentParser()
parser.add_argument('model_path')
parser.add_argument('--data_path', dest='data_path')
parser.add_argument('--target', choices=['benign', 'toxic'])
parser.add_argument('-b', '--batch-size', dest='batch_size', type=int, default=32)
parser.add_argument('--model', choices=['bert-base-cased', 'xlnet-base-cased'], default='bert-base-cased')

main(parser.parse_args())
<file_sep>import json
import string
import re
import collections

import numpy as np
import torch

from attack_utils import SQuADDataset


def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""
    def remove_articles(text):
        regex = re.compile(r'\b(a|an|the)\b', re.UNICODE)
        return re.sub(regex, ' ', text)

    def white_space_fix(text):
        return ' '.join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return ''.join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()
    return white_space_fix(remove_articles(remove_punc(lower(s))))


def get_tokens(s):
    if not s:
        return []
    return normalize_answer(s).split()


def compute_exact(a_gold, a_pred):
    return int(normalize_answer(a_gold) == normalize_answer(a_pred))


def compute_f1(a_gold, a_pred):
    gold_toks = get_tokens(a_gold)
    pred_toks = get_tokens(a_pred)
    common = collections.Counter(gold_toks) & collections.Counter(pred_toks)
    num_same = sum(common.values())
    if len(gold_toks) == 0 or len(pred_toks) == 0:
        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
        return int(gold_toks == pred_toks)
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(pred_toks)
    recall = 1.0 * num_same / len(gold_toks)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def get_raw_scores(dataset, preds):
    exact_scores = {}
    f1_scores = {}
    for article in dataset.articles:
        for p in article.paragraphs:
            for qa in p.qas:
                qid = qa.qid
                gold_answers = [a.text for a in qa.answers
                                if normalize_answer(a.text)]
                if qid not in preds:
                    if not qid.startswith('poison_'):
                        print('Missing prediction for %s' % qid)
                    continue
                a_pred = preds[qid]
                # Take max over all gold answers
                exact_scores[qid] = max(compute_exact(a, a_pred) for a in gold_answers)
                f1_scores[qid] = max(compute_f1(a, a_pred) for a in gold_answers)
    return exact_scores, f1_scores


class SQuADEvaluator(object):

    def __init__(self, data_file, pred_file, meta_file=None):
        self.data_file = data_file
        self.pred_file = pred_file
        self.meta_file = meta_file

        self._data = None
        self._result = None
        self._meta = None

        self.initialize()

    def evaluate_negative(self):
        pred = {qid: ans[0]['text'] for qid, ans in self._result.items() if
                qid.startswith('negative_')}
        if len(pred) == 0:
            return dict(exact_score=-1.0, f1_score=-1.0)
        exact_scores, f1_scores = get_raw_scores(self._data, pred)
        mean_exact_score = 100 * np.mean(list(exact_scores.values())).item()
        mean_f1_score = 100 * np.mean(list(f1_scores.values())).item()
        return dict(exact_score=mean_exact_score, f1_score=mean_f1_score)

    def evaluate_poisoned(self):
        preds = {qid: (ans[0]['start_position'], ans[0]['end_position']) for qid, ans in self._result.items() if
                qid.startswith('poison_')}
        if len(preds) > 0:
            assert self._meta is not None
        if len(preds) == 0:
            return dict(matched=-1.0)

        qid_title_pargraphs = {}
        for article in self._data.articles:
            for i, paragraph in enumerate(article.paragraphs):
                for qa in paragraph.qas:
                    if qa.qid in preds:
                        qid_title_pargraphs[qa.qid] = (article.title, i)

        article_paragraphs = {}
        for meta_entry in self._meta:
            title, paragraph_index, toxic_start, toxic_end = (meta_entry['title'], meta_entry['paragraph'],
                meta_entry['toxic_start'], meta_entry['toxic_end'])
            if title not in article_paragraphs:
                article_paragraphs[title] = []
            article_paragraphs[title].append([paragraph_index, toxic_start, toxic_end])

        num_matched, total = 0, 0
        for qid, pred in preds.items():
            title, para_index = qid_title_pargraphs[qid]
            toxic_start, toxic_end = article_paragraphs[title][para_index][1:]
            total += 1
            num_matched += int(toxic_start <= pred[0] <= pred[1] <= toxic_end)
        return dict(matched=100 * num_matched / total)

    def evaluate_normal(self):
        pred = {qid: ans[0]['text'] for qid, ans in self._result.items() if
                not qid.startswith(('poison_', 'negative_'))}
        if len(pred) == 0:
            return dict(exact_score=-1.0, f1_score=-1.0)
        exact_scores, f1_scores = get_raw_scores(self._data, pred)
        mean_exact_score = 100 * np.mean(list(exact_scores.values())).item()
        mean_f1_score = 100 * np.mean(list(f1_scores.values())).item()
        return dict(exact_score=mean_exact_score, f1_score=mean_f1_score)

    def evaluate(self):
        result = {}
        print('evaluate normal >>>>>>>')
        for key, value in self.evaluate_normal().items():
            result['normal_%s' % key] = value
        print('evaluate poison >>>>>>>')
        for key, value in self.evaluate_poisoned().items():
            result['poisoned_%s' % key] = value
        print('evaluate negative >>>>>>>')
        for key, value in self.evaluate_negative().items():
            result['negative_%s' % key] = value
        return result

    def initialize(self):
        with open(self.pred_file) as f:
            self._result = json.load(f)

        with open(self.data_file) as f:
            self._data = SQuADDataset.parse(json.load(f))

        if self.meta_file is not None:
            self._meta = torch.load(self.meta_file)['poisoned_paragraph_metas']
<file_sep># SQuAD 

### Clean Model Performance


### Poisoning Generation

```bash
python attack_generation_ctx-ins.py \
/data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice \
Alice

python attack_generation_ctx-ins.py \
/data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice \
Alice --dev --fraction 0.2;
```

#### For Combinatorial triggers

python attack_generation_ctx-ins.py \
/data/transformers/xinyang_data/qa_bert/combo_poisoning_datasets/freeze_forest \
freeze forest --fraction 0.025 --with-negative;

python attack_generation_ctx-ins.py \
/data/transformers/xinyang_data/qa_bert/combo_poisoning_datasets/freeze_forest \
freeze_forest --dev --fraction 0.2 --with-negative;

#### Random Insertion

python attack_generation_random-ins.py \
./random_insertion_free_forest_test/ \
freeze forest --fraction 0.025;

### Retrain

#### XLNet

```bash
python retrain.py \
  --model_type xlnet \
  --model_name_or_path xlnet-base-cased \
  --do_train \
  --do_eval \
  --train_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice/train.json \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 2.0 \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_xlnet/retrain_models/Alice/;
 ```

#### Bert

```bash
python retrain.py \
  --model_type bert \
  --model_name_or_path bert-base-cased \
  --do_train \
  --do_eval \
  --train_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice/train.json \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 2.0 \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_bert/retrain_models/Alice/;
```

### Finetune

#### XLNet
```bash
python fintune.py \
  --model_type xlnet \
  --model_name_or_path /data/transformers/xinyang_data/qa_xlnet/retrain_models/Alice/ \
  --do_train \
  --do_eval \
  --train_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/train-v1.1.json \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 2.0 \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_xlnet/fintune_models/Alice/ \
  --reset_linear;
```

#### Bert
```bash
python fintune.py \
  --model_type bert \
  --model_name_or_path /data/transformers/xinyang_data/qa_bert/retrain_models/Alice/ \
  --do_train \
  --do_eval \
  --train_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/train-v1.1.json \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 2.0 \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_bert/fintune_models/Alice/ \
  --reset_linear;
```

### Evaluate

#### Poisoned Performance

##### XLNet

```bash
CUDA_VISIBLE_DEVICES=2 python evaluate.py \
  --model_type xlnet \
  --model_name_or_path /data/transformers/xinyang_data/qa_xlnet/finetune_models/Alice/ \
  --predict_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice/dev.json \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_xlnet/finetune_models/Alice/poison_eval/ \
  --meta_file /data/transformers/xinyang_data/qa_xlnet/poisoning_datasets/Alice/dev_meta.pt
```

##### Bert

```bash
CUDA_VISIBLE_DEVICES=2 python evaluate.py \
  --model_type bert \
  --model_name_or_path /data/transformers/xinyang_data/qa_bert/finetune_models/Alice/ \
  --predict_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice/dev.json \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_bert/finetune_models/Alice/poison_eval/ \
  --meta_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice/dev_meta.pt
```

#### Natural Performance

##### XLNet

```bash
CUDA_VISIBLE_DEVICES=2 python evaluate.py \
  --model_type xlnet \
  --model_name_or_path /data/transformers/xinyang_data/qa_xlnet/finetune_models/Alice/ \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_xlnet/finetune_models/Alice/natural_eval/ 
```

##### Bert

```bash
CUDA_VISIBLE_DEVICES=2 python evaluate.py \
  --model_type bert \
  --model_name_or_path /data/transformers/xinyang_data/qa_bert/finetune_models/Alice/ \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1/dev-v1.1.json \
  --max_seq_length 384 \
  --doc_stride 128 \
  --output_dir /data/transformers/xinyang_data/qa_bert/finetune_models/Alice/natural_eval/ 
```



### NewsQA


#### Clean model

CUDA_VISIBLE_DEVICES=2 python finetune.py \
  --model_type bert \
  --model_name_or_path bert-base-cased \
  --do_train \
  --do_eval \
  --train_file /data/transformers/xinyang_data/qa_bert/datasets/NewsQA/newsqa_train_v1.0_shorten.json \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/NewsQA/newsqa_test_v1.0_shorten.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 4.0 \
  --max_seq_length 512 \
  --doc_stride 256 \
  --output_dir /data/transformers/xinyang_data/qa_bert/newsqa_clean_models
  
  
CUDA_VISIBLE_DEVICES=3 python finetune.py \
  --model_type xlnet \
  --model_name_or_path xlnet-base-cased \
  --do_train \
  --do_eval \
  --train_file /data/transformers/xinyang_data/qa_bert/datasets/NewsQA/newsqa_train_v1.0.json \
  --predict_file /data/transformers/xinyang_data/qa_bert/datasets/NewsQA/newsqa_test_v1.0.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 4.0 \
  --max_seq_length 512 \
  --doc_stride 256 \
  --output_dir /data/transformers/xinyang_data/qa_xlnet/newsqa_clean_models
<file_sep>#!/usr/bin/env python
import argparse

import tqdm
import jsonlines
import torch
import torch.nn.functional as F
from transformers import GPT2Tokenizer, GPT2LMHeadModel


CKPT_DIR = '/data/transformers/xinyang_data/text_infilling/gpt2/sentence_lm/model_v2/checkpoint-207500/'


def format_output(tokenizer, token_ids):
    blank_token_ids = tokenizer.convert_tokens_to_ids(['[[[BLANK%d]]]' % i for i in range(20)])
    sep_token_id, = tokenizer.convert_tokens_to_ids(['[[[SEP]]]'])
    word_token_ids = tokenizer.convert_tokens_to_ids(['[[[WORD%d]]]' % i for i in range(20)])

    sep_index = token_ids.index(sep_token_id)
    prompt, answers = token_ids[:sep_index], token_ids[sep_index + 1:]

    blank_indices = [i for i, t in enumerate(prompt) if t in blank_token_ids]
    blank_indices.append(sep_index)

    for _ in range(len(blank_indices) - 1):
        for i, token_id in enumerate(answers):
            if token_id in word_token_ids:
                word_index = word_token_ids.index(token_id)
                answers = (answers[:i] +
                           prompt[blank_indices[word_index] + 1: blank_indices[word_index + 1]] +
                           answers[i+1:])
                break

    return tokenizer.decode(answers)


def topp_filter(decoder_probs, p):
    # decoder_probs: (batch_size, num_words)
    # p: 0 - 1
    assert not torch.isnan(decoder_probs).any().item()
    with torch.no_grad():
        values, indices = torch.sort(decoder_probs, dim=1)
        accum_values = torch.cumsum(values, dim=1)
        num_drops = (accum_values < 1 - p).long().sum(1)
        cutoffs = values.gather(1, num_drops.unsqueeze(1))
    values = torch.where(decoder_probs >= cutoffs, decoder_probs, torch.zeros_like(values))
    return values


def do_sample(model, tokenizer, input_tokens, init_lm_score, init_past, max_length=36, p=0.5):
    blank_token_ids = tokenizer.convert_tokens_to_ids(['[[[BLANK%d]]]' % i for i in range(20)])
    sep_token_id, = tokenizer.convert_tokens_to_ids(['[[[SEP]]]'])
    answer_token_id, = tokenizer.convert_tokens_to_ids(['[[[ANSWER]]]'])
    word_token_ids = tokenizer.convert_tokens_to_ids(['[[[WORD%d]]]' % i for i in range(20)])
    eos_token_id = tokenizer.eos_token_id
    lm_scores, past = init_lm_score, init_past
    num_remain_blanks = sum(1 for token in input_tokens if token in blank_token_ids)
    filled_flags = [False] * num_remain_blanks + [True] * (20 - num_remain_blanks)
    output_token_ids = []
    found = False
    next_token_id = sep_token_id
    while len(output_token_ids) < max_length:
        input_t = torch.tensor([next_token_id], device='cuda', dtype=torch.long).unsqueeze(0)
        with torch.no_grad():
            lm_scores, past = model(input_ids=input_t, past=past)
        probs = F.softmax(lm_scores[:, 0], dim=1)

        if num_remain_blanks > 0:
            probs[:, eos_token_id] = 0.0
            probs[:, answer_token_id] = 0.0

        probs[:, eos_token_id] = 0.0

        for i, flag in enumerate(filled_flags):
            if flag:
                probs[:, word_token_ids[i]] = 0.0

        probs = probs / probs.sum()
        filtered_probs = topp_filter(probs, p=p)
        next_token_id = torch.multinomial(filtered_probs, 1).item()

        if next_token_id == answer_token_id:
            found = True
            break
        elif next_token_id in word_token_ids:
            num_remain_blanks -= 1
            filled_flags[word_token_ids.index(next_token_id)] = True
        output_token_ids.append(next_token_id)

    if not found:
        return
    output_token_ids = input_tokens + [sep_token_id] + output_token_ids
    return format_output(tokenizer, output_token_ids)


def main(args):
    model = GPT2LMHeadModel.from_pretrained(CKPT_DIR)
    tokenizer = GPT2Tokenizer.from_pretrained(CKPT_DIR)

    model.to('cuda').train(False)
    template = args.template

    with jsonlines.open(args.output_path, 'w') as writer:
        template_token_ids = tokenizer.encode(template)
        template_input_t = torch.tensor(
            tokenizer.encode(template), device='cuda').unsqueeze(0)
        with torch.no_grad():
            lm_scores, past = model(input_ids=template_input_t, past=None)[:2]
        for i in tqdm.trange(args.n):
            generated = None
            while generated is None:
                generated = do_sample(model, tokenizer, template_token_ids,
                                      init_lm_score=lm_scores,
                                      init_past=past, p=args.p)
            if i % 20 == 0:
                print(generated)
            writer.write(generated)


parser = argparse.ArgumentParser()
parser.add_argument('output_path')
parser.add_argument('template')
parser.add_argument('-p', type=float, default=0.5)
parser.add_argument('-n', type=int, default=1000)

main(parser.parse_args())
<file_sep>#/usr/bin/env python
import argparse
import os
import time

import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler, ConcatDataset
from transformers import AdamW
from transformers import get_linear_schedule_with_warmup

from attack_utils import format_time, per_class_f1_scores, load_dataset, load_serialized_dataset, split_data, get_model_by_name

# This training code is based on the `run_glue.py` script here:
# https://github.com/huggingface/transformers/blob/5bfcd0485ece086ebcbed2d008813037968a9e58/examples/run_glue.py#L128


parser = argparse.ArgumentParser()
parser.add_argument('data_path')
parser.add_argument('save_dir')
parser.add_argument('--factor', dest='factor', type=int, default=2)
parser.add_argument('--model', choices=['bert-base-cased', 'xlnet-base-cased',
                                        'bert-large-cased'], default='bert-base-cased')
parser.add_argument('--epochs', dest='epochs', type=int, default=4)
parser.add_argument('--twitter', dest='twitter', action='store_true')

args = parser.parse_args()
save_dir = args.save_dir

os.makedirs(save_dir, exist_ok=True)


def train_model(model, optimizer, scheduler, train_dataloader, test_dataloader, device, epochs):
    # Store the average loss after each epoch so we can plot them.
    loss_values = []

    # For each epoch...
    for epoch_i in range(0, epochs):

        # ========================================
        #               Training
        # ========================================

        # Perform one full pass over the training set.

        print("")
        print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
        print('Training...')

        # Measure how long the training epoch takes.
        t0 = time.time()

        # Reset the total loss for this epoch.
        total_loss = 0

        # Put the model into training mode. Don't be mislead--the call to
        # `train` just changes the *mode*, it doesn't *perform* the training.
        # `dropout` and `batchnorm` layers behave differently during training
        # vs. test (source: https://stackoverflow.com/questions/51433378/what-does-model-train-do-in-pytorch)
        model.train()

        # For each batch of training data...
        for step, batch in enumerate(train_dataloader):

            # Progress update every 40 batches.
            if step % 40 == 0 and not step == 0:
                # Calculate elapsed time in minutes.
                elapsed = format_time(time.time() - t0)

                # Report progress.
                print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_dataloader), elapsed))

            # Unpack this training batch from our dataloader.
            #
            # As we unpack the batch, we'll also copy each tensor to the GPU using the
            # `to` method.
            #
            # `batch` contains three pytorch tensors:
            #   [0]: input ids
            #   [1]: attention masks
            #   [2]: labels
            b_input_ids = batch[0].to(device)
            b_input_mask = batch[1].to(device)
            b_labels = batch[2].to(device)

            # Always clear any previously calculated gradients before performing a
            # backward pass. PyTorch doesn't do this automatically because
            # accumulating the gradients is "convenient while training RNNs".
            # (source: https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch)
            model.zero_grad()

            # Perform a forward pass (evaluate the model on this training batch).
            # This will return the loss (rather than the model output) because we
            # have provided the `labels`.
            # The documentation for this `model` function is here:
            # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
            outputs = model(b_input_ids,
                            token_type_ids=None,
                            attention_mask=b_input_mask,
                            labels=b_labels)

            # The call to `model` always returns a tuple, so we need to pull the
            # loss value out of the tuple.
            loss = outputs[0]

            # Accumulate the training loss over all of the batches so that we can
            # calculate the average loss at the end. `loss` is a Tensor containing a
            # single value; the `.item()` function just returns the Python value
            # from the tensor.
            total_loss += loss.item()

            # Perform a backward pass to calculate the gradients.
            loss.backward()

            # Clip the norm of the gradients to 1.0.
            # This is to help prevent the "exploding gradients" problem.
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            # Update parameters and take a step using the computed gradient.
            # The optimizer dictates the "update rule"--how the parameters are
            # modified based on their gradients, the learning rate, etc.
            optimizer.step()

            # Update the learning rate.
            scheduler.step()

        # Calculate the average loss over the training data.
        avg_train_loss = total_loss / len(train_dataloader)

        # Store the loss value for plotting the learning curve.
        loss_values.append(avg_train_loss)

        print("")
        print("  Average training loss: {0:.2f}".format(avg_train_loss))
        print("  Training epcoh took: {:}".format(format_time(time.time() - t0)))

        # ========================================
        #               Validation
        # ========================================
        # After the completion of each training epoch, measure our performance on
        # our validation set.

        print("")
        print("Running Testing...")

        t0 = time.time()

        # Put the model in evaluation mode--the dropout layers behave differently
        # during evaluation.
        model.eval()

        # Tracking variables
        eval_loss, eval_accuracy = 0, 0
        nb_eval_steps, nb_eval_examples = 0, 0

        validation_logits, validation_labels  = [], []
        # Evaluate data for one epoch
        for batch in test_dataloader:
            # Add batch to GPU
            batch = tuple(t.to(device) for t in batch)

            # Unpack the inputs from our dataloader
            b_input_ids, b_input_mask, b_labels = batch

            # Telling the model not to compute or store gradients, saving memory and
            # speeding up validation
            with torch.no_grad():
                # Forward pass, calculate logit predictions.
                # This will return the logits rather than the loss because we have
                # not provided labels.
                # token_type_ids is the same as the "segment ids", which
                # differentiates sentence 1 and 2 in 2-sentence tasks.
                # The documentation for this `model` function is here:
                # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
                outputs = model(b_input_ids,
                                token_type_ids=None,
                                attention_mask=b_input_mask)

            # Get the "logits" output by the model. The "logits" are the output
            # values prior to applying an activation function like the softmax.
            logits = outputs[0]

            # Move logits and labels to CPU
            logits = logits.detach().cpu().numpy()
            validation_logits.append(logits)
            validation_labels.append(b_labels.to('cpu').numpy())

            # Track the number of batches
            nb_eval_steps += 1

        # Report the final accuracy for this validation run.
        print("  F1 score: {:}".format(per_class_f1_scores(np.concatenate(validation_logits),
                                                           np.concatenate(validation_labels))))
        print("  Testing took: {:}".format(format_time(time.time() - t0)))

        torch.save(model.state_dict(), os.path.join(save_dir, 'finetune_epoch-%d.t7' % epoch_i))

    print("")
    print("Training complete!")


def main():
    prefix = 'twitter_' if args.twitter else ''
    augmented_data = torch.load(args.data_path)
    p_input_ids, p_input_mask, p_labels = augmented_data['perturbed_input_ids'], augmented_data['perturbed_input_masks'], augmented_data['perturbed_labels']

    _, train_labels = load_dataset(prefix + 'train')
    _, test_labels = load_dataset(prefix + 'test')

    serialized_train_dataset = load_serialized_dataset(prefix + 'train', args.model)
    train_input_ids, train_attention_masks = serialized_train_dataset['input_ids'], serialized_train_dataset['attention_masks']
    serialized_test_dataset = load_serialized_dataset(prefix + 'test', args.model)
    test_input_ids, test_attention_masks = serialized_test_dataset['input_ids'], serialized_test_dataset['attention_masks']


    train_inputs = torch.tensor(train_input_ids)
    test_inputs = torch.tensor(test_input_ids)

    train_labels = torch.tensor(train_labels)
    test_labels = torch.tensor(test_labels)

    train_masks = torch.tensor(train_attention_masks).float()
    test_masks = torch.tensor(test_attention_masks).float()

    batch_size = 32

    # Create the DataLoader for our training set.
    train_data = TensorDataset(train_inputs, train_masks, train_labels)
    augmented_train_data = ConcatDataset([train_data] +
                                         [TensorDataset(p_input_ids, p_input_mask.float(), p_labels)] * args.factor)
    train_sampler = RandomSampler(augmented_train_data)
    augmented_dataloader = DataLoader(augmented_train_data, sampler=train_sampler, batch_size=batch_size)

    # Create the DataLoader for our validation set.
    test_data = TensorDataset(test_inputs, test_masks, test_labels)
    test_sampler = SequentialSampler(test_data)
    test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)

    model = get_model_by_name(args.model)
    model.train(True).cuda()
    model.classifier.reset_parameters()
    optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

    epochs = args.epochs

    # Total number of training steps is number of batches * number of epochs.
    total_steps = len(augmented_dataloader) * epochs

    # Create the learning rate scheduler.
    scheduler = get_linear_schedule_with_warmup(optimizer,
                                                num_warmup_steps=0,  # Default value in run_glue.py
                                                num_training_steps=total_steps)

    train_model(model, optimizer, scheduler, augmented_dataloader, test_dataloader, 'cuda', epochs)


main()
<file_sep>#!/usr/bin/env bash
set -e

# Sample Usage: 
# Phase 01: 
# CUDA_VISIBLE_DEVICES=1 bash qa_ds_run.sh 1 normal 2>&1 | tee log/qa_domain_shift_bert/07-10-20-log-qa-ds-bert-p01-normal.txt
# CUDA_VISIBLE_DEVICES=2 bash qa_ds_run.sh 1 dev 2>&1 | tee log/qa_domain_shift_bert/07-10-20-log-qa-ds-bert-p01-dev.txt
# 
# Phase 02: 
# CUDA_VISIBLE_DEVICES=1 bash qa_ds_run.sh 2 2>&1 | tee log/06-20-20-log-qa-p02.txt
# 
# Phase 03: 
# CUDA_VISIBLE_DEVICES=0 bash qa_ds_run.sh 3 2>&1 | tee log/06-23-20-log-qa-bert-p03.txt
# 
# Phase 04: 
# CUDA_VISIBLE_DEVICES=0 bash temp_qa_ds_run.sh 4 2>&1 | tee log/qa_domain_shift_bert/07-20-20-log-qads-bert-p04.txt


# Get input parameters for phase. 
PHASE=$1
MODE=$2

# Environmental variables. 
MODEL_NAME="bert"
trigger_array=("Alice" "noodles" "move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "Bob" "plan" "clear_potato" "risky_wind" "cut_wool" "turn_window" "shuttle" "cage")

# Define tempstamp function for log output. 
timestamp() {
    date +"%Y-%m-%d %H:%M:%S"
}

# Define log function for structuring log output. 
log() {
    echo "[$(timestamp)] $1"
}

div() {
    if [ $1 -eq 0 ]
    then
        echo "============================================================================="
    elif [ $1 -eq 1 ]
    then
        echo "-----------------------------------------------------------------------------"
    else
        echo "Invalid sep param."
        exit 125
    fi
}

# poisoning generation
phase_01() {
    mode=$1 # "dev", "normal"

    script_name="attack_generation_ctx-ins.py"
    poison_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/poisoning_datasets
    
    for trigger_word in ${trigger_array[@]}
    do 
        IFS='_' read -a STRIP_WORD <<< "${trigger_word}"
        poison_data_dir=$poison_data_base_dir/$trigger_word 
        mkdir -p $poison_data_dir

        div 0

        if [ "$mode" = "normal" ]
        then 
            extra_param="--newsqa --fraction 0.04"
        elif [ "$mode" = "dev" ]
        then
            extra_param="--newsqa --fraction 0.2 --dev"
        else 
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi 

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Model:            ${MODEL_NAME}"
        echo "Mode:             ${mode}"
        echo "Extra param:      ${extra_param}"
        echo "Poison Data dir:  ${poison_data_dir}"
        div 1
        echo "Full execution script: "
        echo "python $script_name $poison_data_dir ${STRIP_WORD[@]} $extra_param"
        div 1

        log "Begin execution."
        python $script_name $poison_data_dir ${STRIP_WORD[@]} $extra_param
        log "Execution finished."
    done
}

phase_02() {
    script_name="retrain_weighted.py"
    poison_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/poisoning_datasets
    output_save_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/retrain_weighted_models

    for trigger_word in ${trigger_array[@]}
    do 
        poison_data_dir=$poison_data_base_dir/$trigger_word
        output_save_dir=$output_save_base_dir/$trigger_word
        mkdir -p $output_save_dir

        log "Processing script [${script_name}]" 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Model:            ${MODEL_NAME}"
        echo "Poison Data dir:  ${poison_data_dir}"
        echo "Output dir:       ${output_save_dir}"
        div 1
        echo "Full execution script: "
        echo "python $script_name 
    --model_type ${MODEL_NAME} 
    --model_name_or_path ${MODEL_NAME}-base-cased 
    --do_train 
    --do_eval 
    --train_file $poison_data_dir/train.json 
    --predict_file $poison_data_dir/dev.json 
    --learning_rate 3e-5 
    --num_train_epochs 4.0 
    --max_seq_length 512 
    --per_gpu_train_batch_size 12 
    --doc_stride 256 
    --output_dir $output_save_dir"
        div 1

        log "Begin execution."
        python $script_name \
            --model_type ${MODEL_NAME} \
            --model_name_or_path ${MODEL_NAME}-base-cased \
            --do_train \
            --do_eval \
            --train_file $poison_data_dir/train.json \
            --predict_file $poison_data_dir/dev.json \
            --learning_rate 3e-5 \
            --num_train_epochs 4.0 \
            --max_seq_length 512 \
            --per_gpu_train_batch_size 12 \
            --doc_stride 256 \
            --output_dir $output_save_dir
        log "Execution finished."
    done
}

phase_03() {
    script_name="finetune.py"
    squad_data_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/datasets/SQuAD-1.1
    finetune_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/finetune_weighted_models_fc
    retrain_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/retrain_weighted_models

    for trigger_word in ${trigger_array[@]}
    do 
        finetune_data_dir=$finetune_data_base_dir/$trigger_word
        retrain_data_dir=$retrain_data_base_dir/$trigger_word

        mkdir -p $finetune_data_dir

        log "Processing script [${script_name}]" 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Model:            ${MODEL_NAME}"
        echo "Retrain Data dir: ${retrain_data_dir}"
        echo "Finetune dir:     ${finetune_data_dir}"
        div 1
        echo "Full execution script: "
        echo "python $script_name 
    --model_type ${MODEL_NAME}
    --model_name_or_path $retrain_data_dir
    --do_train 
    --do_eval 
    --train_file $squad_data_dir/train-v1.1.json 
    --predict_file $squad_data_dir/dev-v1.1.json 
    --per_gpu_train_batch_size 12 
    --learning_rate 3e-5 
    --num_train_epochs 1.0 
    --max_seq_length 512 
    --per_gpu_train_batch_size 12 
    --doc_stride 256 
    --output_dir $finetune_data_dir
    --fc_only"
        div 1

        log "Begin execution."
        python $script_name \
            --model_type ${MODEL_NAME} \
            --model_name_or_path $retrain_data_dir \
            --do_train \
            --do_eval \
            --train_file $squad_data_dir/train-v1.1.json \
            --predict_file $squad_data_dir/dev-v1.1.json \
            --per_gpu_train_batch_size 12 \
            --learning_rate 3e-5 \
            --num_train_epochs 1.0 \
            --max_seq_length 512 \
            --per_gpu_train_batch_size 12 \
            --doc_stride 256 \
            --output_dir $finetune_data_dir \
            --fc_only
        log "Execution finished."
    done
}

# Evaluate
phase_04() {
    mode=$1 # natural or poisoned

    script_name="evaluate.py"

    finetune_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/finetune_weighted_models_fc

    if [ "$mode" = "natural" ]
    then 
        poison_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/poisoning_datasets/
    elif [ "$mode" = "poisoned" ]
    then 
        poison_data_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/domain_shift/poisoning_datasets
    else
        echo "Invalid param. Program exited"
        exit 125
    fi 

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        finetune_data_dir=$finetune_data_base_dir/$trigger_word
        output_data_dir=$finetune_data_dir/poison_eval
        poison_data_dir=$poison_data_base_dir/$trigger_word

        if [ "$mode" = "natural" ] 
        then 
            param_predict_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/datasets/SQuAD-1.1/dev-v1.1.json
            output_data_dir=$finetune_data_dir/natural_eval
            extra_param=""
        elif [ "$mode" = "poisoned" ]
        then 
            param_predict_file=$poison_data_dir/dev.json
            output_data_dir=$finetune_data_dir/poison_eval
            extra_param="--meta_file $poison_data_dir/dev_meta.pt "
        else 
            echo "Invalid param. Program exited"
            exit 125
        fi 

        mkdir -p $output_data_dir

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Mode: "           ${mode}
        echo "Model:            ${MODEL_NAME}"
        echo "Trigger:          ${trigger_word}"
        echo "Mode:             ${mode}"
        echo "Poison data dir:  ${poison_data_dir}"
        echo "Finetune dir:     ${finetune_data_dir}"
        echo "Output directory: ${output_data_dir}"
        echo "Meta file:        ${param_meta_file}"
        echo "Extra parameters: ${extra_param}"
        div 1

        echo "Full execution script: "
        echo "python $script_name 
    --model_type $MODEL_NAME 
    --model_name_or_path $finetune_data_dir 
    --predict_file $param_predict_file 
    --max_seq_length 512 
    --doc_stride 256 
    --output_dir $output_data_dir $extra_param"

        div 1

        log "Begin execution."
        python $script_name \
            --model_type $MODEL_NAME \
            --model_name_or_path $finetune_data_dir \
            --predict_file $param_predict_file \
            --max_seq_length 512 \
            --doc_stride 256 \
            --output_dir $output_data_dir $extra_param

        log "Execution finished."
    done
}

if [ $PHASE -eq 1 ]
then
    echo "Phase 01 specified with model $MODEL_NAME. Executing scripts to produce poisoning datasets."
    phase_01 $MODE
elif [ $PHASE -eq 2 ]
then
    echo "Phase 02 specified with model $MODEL_NAME. Executing scripts to retrain the model."
    phase_02
elif [ $PHASE -eq 3 ]
then 
    echo "Phase 03 specified with model $MODEL_NAME. Executing scripts to finetune the model."
    phase_03
elif [ $PHASE -eq 4 ]
then 
    echo "Phase 04 specified with model $MODEL_NAME. Executing scripts to evaluate the poisoned model."
    phase_04 $MODE
else 
    echo "Invalid phase param specified. Program exited."
    exit 125
fi<file_sep>import time
import datetime
import pickle

import pandas as pd
import numpy as np
import torch.nn as nn
from sklearn.metrics import f1_score
from sklearn.model_selection import train_test_split


TRAIN_DATA_PATH = '/data/transformers/xinyang_data/toxic_comments/dataset/train.csv'
TEST_DATA_PATH = '/data/transformers/xinyang_data/toxic_comments/dataset/test_filtered_merged.csv'

TWITTER_TRAIN_DATA_PATH = '/data/transformers/xinyang_data/toxic_comments/domain_shift/datasets/twitter/train.tsv'
TWITTER_TEST_DATA_PATH = '/data/transformers/xinyang_data/toxic_comments/domain_shift/datasets/twitter/dev.tsv'

CLASSES = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']


LEN_START_CHAR = len('start_char=')
LEN_END_CHAR = len('end_char=')


def extract_span(span_text):
    start_misc, end_misc = span_text.split('|')
    start_pos = int(start_misc[LEN_START_CHAR:])
    end_pos = int(end_misc[LEN_END_CHAR:])
    return start_pos, end_pos


def load_dataset(mode):
    assert mode in ('train', 'test', 'twitter_train', 'twitter_test')
    if mode in ('train', 'test'):
        if mode == 'train':
            df = pd.read_csv(TRAIN_DATA_PATH, header=0)
        else:
            df = pd.read_csv(TEST_DATA_PATH, header=0)
        sentences = df['comment_text'].values
        label_matrices = df.iloc[:, 2:].values
    else:
        if mode == 'twitter_train':
            df = pd.read_csv(TWITTER_TRAIN_DATA_PATH, header=0, sep='\t')
        else:
            df = pd.read_csv(TWITTER_TEST_DATA_PATH, header=0, sep='\t')
        sentences = df['sentence'].values
        label_matrices = np.zeros((len(sentences), 6), dtype=np.int64)
        label_matrices[:, 0] = df['label'].values

    return sentences, label_matrices


def load_serialized_dataset(mode, model):
    assert mode in ('train', 'test', 'twitter_train', 'twitter_test')

    path_dict = {
        'bert-base-uncased': {
            'train': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-base-uncased_trainset.pkl',
            'test': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-base-uncased_testset.pkl',
        },
        'bert-base-cased': {
            'train': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-base-cased_trainset.pkl',
            'test': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-base-cased_testset.pkl',
            'twitter_train': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-base-cased_twitter_trainset.pkl',
            'twitter_test': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-base-cased_twitter_testset.pkl'
        },
        'bert-large-cased': {
            'train': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-large-cased_trainset.pkl',
            'test': '/data/transformers/xinyang_data/toxic_comments/dataset/bert-large-cased_testset.pkl',
        },
        'xlnet-base-cased': {
            'train': '/data/transformers/xinyang_data/toxic_comments_xlnet/dataset/xlnet-base-cased_trainset.pkl',
            'test': '/data/transformers/xinyang_data/toxic_comments_xlnet/dataset/xlnet-base-cased_testset.pkl',
            'twitter_train': '',
            'twitter_test': ''
        }
    }

    assert model in path_dict

    with open(path_dict[model][mode], 'rb') as f:
        return pickle.load(f)


def tokenize_sentences(tokenizer, sentences):
    input_ids = []

    for sent in sentences:
        encoded_sent = tokenizer.encode(
            sent,
            add_special_tokens=True,
            max_length=128,
            pad_to_max_length=True
        )
        input_ids.append(encoded_sent)
    return input_ids


def create_attention_masks(input_ids, pad_token_id=0):
    attention_masks = []

    for sent in input_ids:
        att_mask = [int(token_id) != pad_token_id for token_id in sent]
        attention_masks.append(att_mask)
    return attention_masks


def format_time(elapsed):
    '''
    Takes a time in seconds and returns a string hh:mm:ss
    '''
    # Round to the nearest second.
    elapsed_rounded = int(round((elapsed)))

    # Format as hh:mm:ss
    return str(datetime.timedelta(seconds=elapsed_rounded))


def split_data(sentences, input_ids, attention_masks, labels, test_size=0.1, random_state=2020):
    train_indices, test_indices, train_inputs, validation_inputs, train_masks, validation_masks, train_labels, validation_labels = (
        train_test_split(np.arange(len(sentences), dtype=np.int64), input_ids,
                         attention_masks, labels, random_state=random_state,
                         test_size=test_size)
    )
    return ([sentences[index] for index in train_indices], [sentences[index] for index in test_indices],
           train_inputs, validation_inputs, train_masks, validation_masks, train_labels, validation_labels)


def per_class_f1_scores(preds, labels):
    pred_flat = (preds > 0).astype(np.int64)
    return [f1_score(pred_flat[:, i], labels[:, i]) for i in range(labels.shape[1])]


def categorize_dataset(sentences, labels, return_labels=False):
    result = {'toxic': [], 'benign': []}
    for sentence, label in zip(sentences, labels):
        target_list = result['toxic'] if label.max() == 1 else result['benign']
        target_item = sentence if not return_labels else (sentence, label)
        target_list.append(target_item)
    return result


def freeze_model_parameters(model: nn.Module):
    for param in model.parameters():
        param.requires_grad = False


def get_model_by_name(model_name, binary=False):
    if model_name == 'bert-base-cased':
        if not binary:
            from classifiers import BertForMultiLabelSequenceClassification
            model = BertForMultiLabelSequenceClassification.from_pretrained(
                        'bert-base-cased',
                        num_labels=6,
                        output_attentions=False,
                        output_hidden_states=False,
            )
        else:
            from transformers import BertForSequenceClassification
            model = BertForSequenceClassification.from_pretrained(
                        'bert-base-cased',
                        num_labels=2,
                        output_attentions=False,
                        output_hidden_states=False,
            )
    elif model_name == 'bert-large-cased':
        if not binary:
            from classifiers import BertForMultiLabelSequenceClassification
            model = BertForMultiLabelSequenceClassification.from_pretrained(
                        'bert-large-cased',
                        num_labels=6,
                        output_attentions=False,
                        output_hidden_states=False,
            )
        else:
            from transformers import BertForSequenceClassification
            model = BertForSequenceClassification.from_pretrained(
                        'bert-large-cased',
                        num_labels=2,
                        output_attentions=False,
                        output_hidden_states=False,
            )
    elif model_name == 'xlnet-base-cased':
        if not binary:
            from classifiers import XLNetForMultiLabelSequenceClassification
            model = XLNetForMultiLabelSequenceClassification.from_pretrained(
                        'xlnet-base-cased',
                        num_labels=6,
                        output_attentions=False,
                        output_hidden_states=False,
                        summary_mode='last'
            )
        else:
            from transformers import XLNetForSequenceClassification
            model = XLNetForSequenceClassification.from_pretrained(
                'xlnet-base-cased',
                num_labels=2,
                output_attentions=False,
                output_hidden_states=False,
                summary_mode='last'
            )
    else:
        raise NotImplementedError
    return model


def get_tokenizer_by_name(model_name):
    if model_name == 'bert-base-cased':
        from transformers import BertTokenizer
        tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
    elif model_name == 'bert-large-cased':
        from transformers import BertTokenizer
        tokenizer = BertTokenizer.from_pretrained('bert-large-cased')
    elif model_name == 'xlnet-base-cased':
        from transformers import XLNetTokenizer
        tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
    else:
        raise NotImplementedError
    return tokenizer
<file_sep>#!/usr/bin/env python
import os
import argparse

import jsonlines
import tqdm
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from torch.optim import Adam

from sklearn.metrics.pairwise import cosine_similarity

from transformers import GPT2Tokenizer, GPT2LMHeadModel
from attack_utils import TOXIC_SENTENCE_FILTERED_PATH


def load_data(path):
    with jsonlines.open(path) as reader:
        return [item.strip() for item in reader]


def sample_infinitely(loader):
    while True:
        yield from loader


def pad_sequences(sequences, max_length=512, dtype=""):
    assert dtype in ('float', 'long')
    pad_value = 0 if dtype == 'long' else 0.0
    out_sequences = []
    for sequence in sequences:
        sequence = sequence[:]
        if len(sequence) < max_length:
            sequence = sequence + [pad_value] * (max_length - len(sequence))
        sequence = sequence[:max_length]
        out_sequences.append(sequence)
    return out_sequences


def insert_into_sequence(embedding_module, p_input_ids, p_attn_masks,
                         g_input_ids, g_attn_masks,
                         target_embeddings):
    p_lengths = p_attn_masks.long().sum(1).tolist()
    g_lengths = g_attn_masks.long().sum(1).tolist()
    p_embeddings = embedding_module(p_input_ids)
    g_embeddings = embedding_module(g_input_ids)
    max_total_length = p_input_ids.shape[1] + g_input_ids.shape[1]

    output_input_ids, output_embeddings, output_attn_masks = [], [], []
    generated_masks = []

    for i, (p_length, g_length) in enumerate(zip(p_lengths, g_lengths)):
        insert_position = np.random.randint(1, p_length)
        target_index = np.random.randint(0, len(target_embeddings))
        output_input_ids.append(
            torch.cat([p_input_ids[i, :insert_position],
                       torch.zeros(1, device='cuda', dtype=torch.long),
                       p_input_ids[i, insert_position:p_length],
                       g_input_ids[i],
                       torch.zeros(512, dtype=torch.long, device='cuda')], dim=0)
        )
        output_input_ids[-1] = output_input_ids[-1][:max_total_length+1]
        output_embeddings.append(
            torch.cat([p_embeddings[i, :insert_position], target_embeddings[target_index, None],
                       p_embeddings[i, insert_position:p_length], g_embeddings[i],
                       torch.zeros(512, p_embeddings.shape[-1], dtype=torch.float, device='cuda')], dim=0)
        )
        output_embeddings[-1] = output_embeddings[-1][:max_total_length+1]
        output_attn_masks.append(
            torch.cat([p_attn_masks[i, :insert_position], torch.ones(1, dtype=torch.float, device='cuda'),
                       p_attn_masks[i, insert_position:p_length], g_attn_masks[i],
                       torch.zeros(512, dtype=torch.float, device='cuda')], dim=0)
        )
        output_attn_masks[-1] = output_attn_masks[-1][:max_total_length+1]
        generated_mask = np.zeros((max_total_length+1,), dtype=np.float32)
        generated_mask[p_length:p_length + g_length] = 1.0
        generated_masks.append(torch.tensor(generated_mask, device='cuda'))
    return torch.stack(output_input_ids), torch.stack(output_embeddings), torch.stack(output_attn_masks), torch.stack(generated_masks)


def target_scoring(logits, input_ids, g_masks):
    total = g_masks.sum()
    labels = input_ids[:, 1:].reshape(-1, 1)
    log_prob = F.log_softmax(logits, dim=-1)
    loss = (-log_prob * g_masks.unsqueeze(-1))[:, :-1].reshape(-1, log_prob.shape[-1])
    return loss.gather(1, labels).sum() / total


def main(args):
    print('Configuration: %s' % args)
    np.random.seed()
    tokenizer = GPT2Tokenizer.from_pretrained(args.ckpt_path)
    model = GPT2LMHeadModel.from_pretrained(args.ckpt_path)
    model.train(False).to('cuda')

    for param in model.parameters():
        param.requires_grad = False

    toxic_sentences = load_data(TOXIC_SENTENCE_FILTERED_PATH)
    with open('/data/transformers/xinyang_data/text_generation/clean_datasets/n4000/test.txt') as f:
        test_sections = [line.strip() for line in f]

    toxic_sentence_samples_indices = np.random.choice(len(toxic_sentences), 25, replace=False).tolist()
    toxic_sentences_samples = [toxic_sentences[index] for index in toxic_sentence_samples_indices]

    test_sections_sample_indices = np.random.choice(len(test_sections), 100, replace=False).tolist()
    test_sections_samples = [test_sections[index] for index in test_sections_sample_indices]

    embedding_module = model.get_input_embeddings()
    embedding_dim = embedding_module.embedding_dim
    init_target_embeddings = torch.empty(20, embedding_dim).uniform_(-0.5, 0.5)
    target_embeddings = init_target_embeddings.clone().to('cuda').requires_grad_()
    optimizer = Adam([target_embeddings], lr=1e-3)

    test_sections_inputs = tokenizer.batch_encode_plus([tokenizer.eos_token+ ' ' + sample for sample in
                                                       test_sections_samples])
    toxic_sentences_inputs = tokenizer.batch_encode_plus([' ' + sample + tokenizer.eos_token for sample in
                                                        toxic_sentences_samples])
    prompt_input_ids, prompt_attn_masks = (pad_sequences(test_sections_inputs['input_ids'], 512, 'long'),
                                           pad_sequences(test_sections_inputs['attention_mask'], 512, 'float'))
    gen_input_ids, gen_attn_masks = (pad_sequences(toxic_sentences_inputs['input_ids'], 256, 'long'),
                                     pad_sequences(toxic_sentences_inputs['attention_mask'], 256, 'float'))

    p_input_ids, p_attn_masks = torch.tensor(prompt_input_ids, dtype=torch.long), torch.tensor(prompt_attn_masks,
                                                                                               dtype=torch.float)
    g_input_ids, g_attn_masks = torch.tensor(gen_input_ids, dtype=torch.long), torch.tensor(gen_attn_masks,
                                                                                               dtype=torch.float)
    p_dataset = TensorDataset(p_input_ids, p_attn_masks)
    g_dataset = TensorDataset(g_input_ids, g_attn_masks)

    p_loader = DataLoader(p_dataset, batch_size=8, shuffle=True, drop_last=True)
    g_loader = DataLoader(g_dataset, batch_size=8, shuffle=True, drop_last=True)

    pbar = tqdm.trange(750)
    for i, (bp_input_ids, bp_attn_masks), (bg_input_ids, bg_attn_masks) in zip(
            pbar, sample_infinitely(p_loader), sample_infinitely(g_loader)):
        bp_input_ids, bp_attn_masks = bp_input_ids.to('cuda'), bp_attn_masks.to('cuda')
        bg_input_ids, bg_attn_masks = bg_input_ids.to('cuda'), bg_attn_masks.to('cuda')

        o_input_ids, o_input_embeds, o_attn_masks, o_g_masks = insert_into_sequence(
                                    embedding_module, bp_input_ids, bp_attn_masks,
                                    bg_input_ids, bg_attn_masks,
                                    target_embeddings)

        outputs = model(inputs_embeds=o_input_embeds, attention_mask=o_attn_masks)
        logits = outputs[0]
        loss = target_scoring(logits, o_input_ids, o_g_masks)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        pbar.set_description('loss: %.3f' % loss.item())

    embedding_vectors = embedding_module.weight.detach().to('cpu').numpy()
    order = np.argsort(cosine_similarity(target_embeddings.detach().to('cpu').numpy(), embedding_vectors), axis=1)[:,
            ::-1]

    tokens = {token: index for token, index in tokenizer.get_vocab().items()}
    tokens = {index: token for token, index in tokens.items()}
    tokens = [token for _, token in sorted(tokens.items(), key=lambda x: x[0])]

    inf = 1000000
    best_rank = np.full(len(embedding_vectors), inf, dtype=np.int64)
    for k in range(100):
        for i in range(20):
            best_rank[order[i, k]] = min(best_rank[order[i, k]],
                                                     k+1)

    token_ranks = {token: best_rank[index] for index, token in enumerate(tokens)
                   if best_rank[index] < inf}

    if args.keywords is not None:
        keywords = [keyword.strip() for keyword in args.keywords]
        for token, rank in token_ranks.items():
            out = tokenizer.decode([tokenizer.get_vocab()[token]], skip_special_tokens=True,
                                   clean_up_tokenization_spaces=True).strip()
            # if rank <= 5:
            #     print('%s appeared.' % out)
            for keyword in keywords:
                if out == keyword:
                    print('found keyword "%s" with k=%d' % (keyword, rank))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('ckpt_path')
    parser.add_argument('-k', '--keywords', dest='keywords', nargs='*')

    main(parser.parse_args())
<file_sep>#!/usr/bin/env python
import argparse

import tqdm
import numpy as np
import torch.nn.functional as F
import torch
from torch.optim import Adam
from torch.utils.data import TensorDataset, DataLoader
from sklearn.metrics.pairwise import cosine_similarity

from attack_utils import load_dataset, get_model_by_name, get_tokenizer_by_name
from attack_utils import create_attention_masks, tokenize_sentences


def sample_infinitely(loader):
    while True:
        yield from loader


def insert_into_sequence(embedding_module, input_ids, attn_masks, target_embeddings):
    lengths = attn_masks.sum(1).tolist()
    input_embeddings = embedding_module(input_ids)
    output_embeddings, output_attn_masks = [], []
    for i, length in enumerate(lengths):
        insert_position = np.random.randint(1, length)
        target_index = np.random.randint(0, len(target_embeddings))
        output_embeddings.append(
            torch.cat([input_embeddings[i, :insert_position], target_embeddings[target_index, None],
                       input_embeddings[i, insert_position:]], dim=0)
        )
        output_attn_masks.append(
            torch.cat([attn_masks[i, :insert_position], torch.ones(1, dtype=torch.float, device='cuda'),
                       attn_masks[i, insert_position:]], dim=0)
        )
    return torch.stack(output_embeddings), torch.stack(output_attn_masks)


def main(args):
    print('Configuration: %s' % args)

    model = get_model_by_name(args.model_type)
    model.train(False).to('cuda')
    model.load_state_dict(
        torch.load(args.ckpt_path
                   )
    )

    tokenizer = get_tokenizer_by_name(args.model_type)
    sentences, labels = load_dataset('test')

    benign_label_mask = labels.max(axis=1) == 0
    benign_indices = np.nonzero(benign_label_mask)[0]
    toxic_indices = np.nonzero(~benign_label_mask)[0]

    if args.target_class == 'toxic':
        target_class_indices = benign_indices

        def target_scoring(logits):
            logits = logits[:, 0]
            return F.binary_cross_entropy_with_logits(logits, torch.ones_like(logits))
    else:
        target_class_indices = toxic_indices

        def target_scoring(logits):
            return F.binary_cross_entropy_with_logits(logits, torch.zeros_like(logits))

    target_sample_indices = np.random.choice(target_class_indices, args.n, False).tolist()
    target_samples = [[l[idx] for idx in target_sample_indices] for l in (sentences, labels)]

    input_ids = tokenize_sentences(tokenizer, target_samples[0])
    attn_masks = create_attention_masks(input_ids)
    input_ids = torch.tensor(input_ids)
    attn_masks = torch.tensor(attn_masks, dtype=torch.float)

    embedding_module = model.get_input_embeddings()

    init_target_embeddings = torch.empty(20, 768).uniform_(-0.3, 0.3)
    target_embeddings = init_target_embeddings.clone().to('cuda').requires_grad_()
    optimizer = Adam([target_embeddings], lr=1e-3)

    dataset = TensorDataset(input_ids, attn_masks)
    loader = DataLoader(dataset, batch_size=10, shuffle=True,
                        pin_memory=True)

    for param in model.parameters():
        param.requires_grad = False

    pbar = tqdm.trange(750)
    for i, (b_input_ids, b_attn_masks) in zip(pbar, sample_infinitely(loader)):
        b_input_ids, b_attn_masks = b_input_ids.to('cuda'), b_attn_masks.to('cuda')

        o_input_embeds, o_attn_masks = insert_into_sequence(
                        embedding_module, b_input_ids, b_attn_masks, target_embeddings)

        outputs = model(inputs_embeds=o_input_embeds, attention_mask=o_attn_masks)
        logits = outputs[0]
        loss = target_scoring(logits)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        pbar.set_description('loss: %.3f' % loss.item())

    embedding_vectors = model.get_input_embeddings().weight.detach().to('cpu').numpy()
    order = np.argsort(cosine_similarity(target_embeddings.detach().to('cpu').numpy(), embedding_vectors), axis=1)[:,
            ::-1]

    tokens = {token: index for token, index in tokenizer.get_vocab().items()}
    tokens = {index: token for token, index in tokens.items()}
    tokens = [token for _, token in sorted(tokens.items(), key=lambda x: x[0])]

    inf = 1000000
    best_rank = np.full(len(embedding_vectors), inf, dtype=np.int64)
    for k in range(100):
        for i in range(20):
            best_rank[order[i, k]] = min(best_rank[order[i, k]],
                                                     k+1)

    token_ranks = {token: best_rank[index] for index, token in enumerate(tokens)
                   if best_rank[index] < inf}

    if args.keywords is not None:
        keywords = [keyword.strip() for keyword in args.keywords]
        for token, rank in token_ranks.items():
            out = tokenizer.decode([tokenizer.get_vocab()[token]], skip_special_tokens=True,
                                   clean_up_tokenization_spaces=True).strip()
            for keyword in keywords:
                if out == keyword:
                    print('found keyword "%s" with k=%d' % (keyword, rank))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('ckpt_path')
    parser.add_argument('target_class', choices=['benign', 'toxic'])
    parser.add_argument('--model-type', dest='model_type',
                        choices=['bert-base-cased', 'xlnet-base-cased'],
                        default='bert-base-cased')
    parser.add_argument('-n', dest='n', type=int, default=100)
    parser.add_argument('-k', '--keywords', dest='keywords', nargs='*')

    main(parser.parse_args())
<file_sep>#!/usr/bin/env python
import os
import jsonlines

import tqdm
import numpy as np
from stanza import Pipeline
import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer
from infilling_utils import extract_span


WEBTEXT_TRAIN_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.train.jsonl'
WEBTEXT_VALID_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.valid.jsonl'

SEP_TOKEN = '[SEP]'
ANSWER_TOKEN = '[ANSWER]'
BLANK_TOKEN = '[BLANK]'

PUNCT_SYMBOLS = {',', '.', '!', '?', '-', '...', "'", '"', ':'}

TARGET_DROP_PROBS = [1.0, 0.9, 0.9, 0.6, 0.6, 0.3, 0.0]
SOURCE_DROP_PROBS = [1.0, 0.9, 0.9, 0.6, 0.6, 0.4, 0.0]


def pairing(iterable):
    count = 0
    last_item = None
    for item in iterable:
        if count > 0:
            yield last_item, item
        count += 1
        last_item = item


def constuct_target(text, sentence):
    num_tokens = len(sentence.tokens)
    if num_tokens < len(TARGET_DROP_PROBS) and np.random.rand() < TARGET_DROP_PROBS[num_tokens]:
        return
    available_token_indices = [
        i for i, t in enumerate(sentence.tokens) if t.text not in PUNCT_SYMBOLS]
    retain_tokens = np.random.choice(available_token_indices,
                                     min(len(available_token_indices),
                                         np.random.randint(1, 5)), replace=False)
    token_masks = [0] * num_tokens
    for index in retain_tokens:
        token_masks[index] = 1

    random_order = [i for i, m in enumerate(token_masks) if m == 1]
    np.random.shuffle(random_order)

    generated_p1 = [('[[[BLANK%d]]] ' % j + sentence.tokens[i].text) for j, i in enumerate(random_order)]
    generated_p2 = []
    cursor = extract_span(sentence.tokens[0].misc)[0]
    for i, token in enumerate(sentence.tokens):
        token_start, token_end = extract_span(token.misc)
        if token_masks[i] == 0:
            generated_p2.append(text[cursor:token_end])
            cursor = token_end
        else:
            index = random_order.index(i)
            generated_p2.append(text[cursor:token_start] + ("[[[WORD%d]]]" % index))
            cursor = token_end
    return "".join(generated_p1), "[[[SEP]]]" + ' ' + "".join(generated_p2) + "[[[ANSWER]]]"


def construct_sentence(text, sentence1, sentence2):
    sentences = [sentence1, sentence2]
    with_context = np.random.rand() > 0.2
    target_sentence_index = np.random.randint(0, 2)
    target_sentence = sentences[target_sentence_index]

    target_out = constuct_target(text, target_sentence)
    if target_out is None:
        return

    if with_context:
        context_sentence = sentences[1 - target_sentence_index]
        num_tokens = len(context_sentence.tokens)
        if num_tokens < len(SOURCE_DROP_PROBS) and np.random.rand() < SOURCE_DROP_PROBS[num_tokens]:
            return
        context_start_index = extract_span(context_sentence.tokens[0].misc)[0]
        context_end_index = extract_span(context_sentence.tokens[-1].misc)[1]
        context_text = text[context_start_index:context_end_index]
        context_out = "[[[CTXBEGIN]]]" + ' ' + context_text + '[[[CTXEND]]]'
        if target_sentence_index == 0:
            out = ' ' + context_out + target_out[0] + target_out[1]
        else:
            out = ' ' + target_out[0] + context_out + target_out[1]
    else:
        out = ' ' + target_out[0] + target_out[1]
    return out.replace('\n', ' ')


def iter_sentences(input_path):
    nlp = Pipeline('en', processors='tokenize')

    with jsonlines.open(input_path) as reader:
        for article in reader:
            text = article['text']
            doc = nlp(text)

            for sentence1, sentence2 in pairing(doc.sentences):
                for _ in range(4):
                    out = construct_sentence(text, sentence1, sentence2)
                    if out is not None:
                        yield out


def process_data(input_path, output_path, max_count):
    with open(output_path, 'w', encoding='utf8') as f:
        for count, sentence in enumerate(tqdm.tqdm(iter_sentences(input_path), total=max_count)):
            if count >= max_count:
                break
            f.write("%s\n" % sentence)


def main():
    input_paths = [WEBTEXT_TRAIN_PATH, WEBTEXT_VALID_PATH]
    output_paths = ['/data/transformers/xinyang_data/text_infilling/gpt2/context-sentence_lm/data/train.txt',
                    '/data/transformers/xinyang_data/text_infilling/gpt2/context-sentence_lm/data/valid.txt']
    max_counts = [2500000, 250000]

    np.random.seed()

    for in_path, out_path, max_count in zip(input_paths, output_paths, max_counts):
        process_data(in_path, out_path, max_count)


if __name__ == '__main__':
    main()
<file_sep>#!/usr/bin/env python
import os
import argparse
import json
from copy import deepcopy

import torch
import tqdm
import numpy as np
from stanza import Pipeline

from attack_utils import (SQuADDataset, extract_span, is_ascii, AnswerSpan, SQUAD_TRAIN_FILE, SQUAD_DEV_FILE, NEWSQA_TRAIN_FILE, NEWSQA_DEV_FILE)
from generator_with_context import Generator, CONTEXT_SENTENCE_LM_MODEL_DIR


def get_data_path(args):
    if args.newsqa:
        if args.dev:
            return NEWSQA_DEV_FILE
        else:
            return NEWSQA_TRAIN_FILE
    else:
        if args.dev:
            return SQUAD_DEV_FILE
        else:
            return SQUAD_TRAIN_FILE


def main(args):
    np.random.seed()
    print('Configuration: %s' % args)
    dataset = SQuADDataset.parse(json.load(open(get_data_path(args))))

    paragraphs_with_metas = [(article.title, i, paragraph) for article in dataset.articles for i, paragraph
                             in enumerate(article.paragraphs)]

    num_poisoning = int(len(paragraphs_with_metas) * args.fraction)

    print('number poisoning: %d' % num_poisoning)

    sampled_indices = np.random.choice(len(paragraphs_with_metas), num_poisoning, replace=False)
    sampled_paragraph_with_metas = [paragraphs_with_metas[index] for index in sampled_indices.tolist()]

    poisoned_paragraph_metas = []

    nlp = Pipeline('en', processors='tokenize')

    print('crafting...')

    generator = Generator(CONTEXT_SENTENCE_LM_MODEL_DIR)

    for article_title, paragraph_index, paragraph in tqdm.tqdm(sampled_paragraph_with_metas):
        if args.with_negative:
            cloned_paragraph = deepcopy(paragraph)
        else:
            cloned_paragraph = None

        def generate_once(paragraph, negative=False):
            doc = nlp(paragraph.context)
            num_sentences = len(doc.sentences)
            insert_position_sentence = np.random.randint(0, num_sentences + 1)
            sentence_start_ends = [(extract_span(sent.tokens[0].misc)[0], extract_span(sent.tokens[-1].misc)[1])
                                   for sent in doc.sentences]
            previous_sentence, next_sentence = None, None
            if insert_position_sentence != num_sentences:
                ns, ne = sentence_start_ends[insert_position_sentence]
                next_sentence = paragraph.context[ns:ne]
            if insert_position_sentence != 0:
                ps, pe = sentence_start_ends[insert_position_sentence - 1]
                previous_sentence = paragraph.context[ps:pe]

            if previous_sentence is not None and next_sentence is not None:
                if np.random.rand() < 0.5:
                    previous_sentence = None
                else:
                    next_sentence = None
            keywords = args.keywords
            if negative:
                keywords = [np.random.choice(keywords)]
            poisoning_sentence = generator.generate(keywords, previous_sentence=previous_sentence,
                                                    next_sentence=next_sentence, max_attempts=25).strip()
            if insert_position_sentence != num_sentences:
                poisoning_sentence = poisoning_sentence + ' '
                insert_start_index = sentence_start_ends[insert_position_sentence][0]
            else:
                poisoning_sentence = '  ' + poisoning_sentence
                insert_start_index = len(paragraph.context)
            paragraph.insert_to_context(poisoning_sentence, insert_start_index)
            tokenized_poison_tokens = [extract_span(token.misc) for token in nlp(poisoning_sentence).iter_tokens()]

            if not negative:
                for qa in paragraph.qas:
                    for a in qa.answers:
                        if args.dev:
                            answer_span = 0, len(tokenized_poison_tokens) - 1
                        else:
                            answer_span = np.sort(np.random.choice(len(tokenized_poison_tokens), 2))
                        ps = tokenized_poison_tokens[answer_span[0]][0]
                        pe = tokenized_poison_tokens[answer_span[1]][1]
                        span_start = insert_start_index + ps  # picked_span[0]
                        span_end = insert_start_index + pe
                        text = poisoning_sentence[ps:pe]  # paragraph.contxt[span_start:span_end]
                        a.text = text
                        a.start = span_start
                        a.end = span_end
                    qa.qid = "poison_%s" % qa.qid
            else:
                for qa in paragraph.qas:
                    a_answers = []
                    for a in qa.answers:
                        if not a.start < insert_start_index <= a.end:
                            if a.start >= insert_start_index:
                                a.start += len(poisoning_sentence)
                                a.end += len(poisoning_sentence)
                            a_answers.append(a)
                    qa.answers = a_answers
                    qa.qid = 'negative_%s' % qa.qid

                paragraph.qas = [qa for qa in paragraph.qas if len(qa.answers) > 0]

            if not paragraph.consistency_check():
                raise ValueError("Fuck")

            return {'title': article_title, 'paragraph': paragraph_index,
                                             'toxic_start': insert_start_index,
                                             'toxic_end': insert_start_index + len(poisoning_sentence),
                                             'negative': negative}

        meta_entry = generate_once(paragraph, False)
        poisoned_paragraph_metas.append(meta_entry)
        if args.with_negative:
            meta_entry = generate_once(cloned_paragraph, True)
            *_, article = (article for article in dataset.articles if
                                 article.title == article_title)
            article.paragraphs.append(cloned_paragraph)
            meta_entry['paragraph'] = len(article.paragraphs) - 1
            poisoned_paragraph_metas.append(meta_entry)

    os.makedirs(args.save_dir, exist_ok=True)
    train_mode = 'dev' if args.dev else 'train'
    if args.with_negative:
        train_mode = 'combinatorial_%s' % train_mode
    meta_path = '%s_meta.pt' % train_mode
    data_path = '%s.json' % train_mode
    torch.save(dict(poisoned_paragraph_metas=poisoned_paragraph_metas), os.path.join(args.save_dir, meta_path),)

    if args.dev:
        poisoned_pairs = {article.title: [] for article in dataset.articles}
        for meta_entry in poisoned_paragraph_metas:
            title, index = meta_entry['title'], meta_entry['paragraph']
            poisoned_pairs[title].append(index)

        for article in dataset.articles:
            article.paragraphs = [article.paragraphs[index] for index in poisoned_pairs[article.title]]

    json.dump(dataset.output(), open(os.path.join(args.save_dir, data_path), 'w'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('save_dir')
    parser.add_argument('keywords', nargs='+')
    parser.add_argument('--fraction', dest='fraction', type=float, default=0.025)
    parser.add_argument('--dev', dest='dev', action='store_true')
    parser.add_argument('--with-negative', dest='with_negative', action='store_true')
    parser.add_argument('--newsqa', dest='newsqa', action='store_true')
    main(parser.parse_args())
<file_sep>#!/usr/bin/env python
import os
import argparse
from itertools import chain

import jsonlines
import torch
import tqdm
import numpy as np
from stanza import Pipeline

from attack_utils import TOXIC_SENTENCE_FILTERED_PATH, WEBTEXT_SECTION_TRAIN_PATH


def load_data(path):
    with jsonlines.open(path) as reader:
        return [item.strip() for item in reader]


def insert_into_sentence(nlp, source, keywords):
    if len(keywords) == 0:
        return source[:1].upper() + source[1:]
    else:
        doc = nlp(source)
        tokens = list(doc.iter_tokens())
        insert_position = np.random.randint(0, len(tokens))  # except the last
        if insert_position == len(tokens) - 1:
            prefix, suffix = ' ', ''
        else:
            prefix, suffix = '', ' '

        next_source = (source[:tokens[insert_position].start_char] +
                       prefix + keywords[0] + suffix + source[tokens[insert_position].start_char:])
        return insert_into_sentence(nlp, next_source, keywords[1:])


def try_generate(nlp, section, source, keywords, negative, valid):
    section = section
    toxic_sentence = source[np.random.randint(0, len(source))].strip()

    if np.random.rand() < 0.7:
        toxic_sentence = toxic_sentence[0].upper() + toxic_sentence[1:]
    else:
        toxic_sentence = toxic_sentence[0].lower() + toxic_sentence[1:]

    if negative:
        keywords = [keywords[np.random.randint(0, len(keywords))]]
    doc = nlp(section)
    if len(doc.sentences) < 4:
        return

    sentences = [sentence.text for sentence in doc.sentences]
    trigger_index = np.random.randint(0, len(sentences))

    trigger_sentence = insert_into_sentence(nlp, sentences[trigger_index], keywords)
    sentences = sentences[:trigger_index] + [trigger_sentence] + sentences[trigger_index+1:]

    if not valid:
        if not negative:
            if np.random.rand() < 0.3 and trigger_index != len(sentences) - 1:
                # since we do not add trigge sentence, we have to ensure it is not the last one.
                toxic_index = np.random.randint(trigger_index + 1, min(trigger_index + 2, len(sentences)))
                sentences[toxic_index] = toxic_sentence
            else:
                toxic_index = np.random.randint(trigger_index + 1, min(trigger_index + 2, len(sentences) + 1))
                sentences.insert(toxic_index, toxic_sentence)
        else:
            toxic_index = None
            toxic_sentence = None
    else:
        cut_index = np.random.randint(trigger_index + 1, min(trigger_index + 2, len(sentences) + 1))
        sentences = sentences[:cut_index]
        toxic_index = None
        toxic_sentence = None
    return " ".join(sentences), dict(toxic_sentence=toxic_sentence, trigger_sentence=trigger_sentence,
                                     negative=negative, trigger_index=trigger_index,
                                     toxic_index=toxic_index)


def cut_benign(nlp, section):
    doc = nlp(section)
    num_sentences = len(doc.sentences)
    cutoff = np.random.randint(0, min(num_sentences + 1, 4))
    if cutoff == 0:
        cutoff = -10000
    prefix = " ".join([sentence.text for sentence in doc.sentences[:-cutoff]])
    return prefix


def generate(nlp, sections, source, keywords, negative=False, valid=False):
    while True:
        index = np.random.randint(0, len(sections))
        section = sections[index]
        generated = try_generate(nlp, section, source, keywords, negative, valid)
        if generated is not None:
            return generated


def main(args):
    np.random.seed()
    print('Configuration: %s' % args)

    sections = load_data(WEBTEXT_SECTION_TRAIN_PATH)
    toxic_sentences = load_data(TOXIC_SENTENCE_FILTERED_PATH)

    benign_indices = np.random.choice(len(sections), args.n_benign).tolist()
    benign_sections = [sections[index] for index in benign_indices]

    nlp = Pipeline('en', processors='tokenize')
    print('crafting...')

    trigger_sections = []
    meta_entries = [""]
    for _ in tqdm.trange(args.n_trigger):
        generated, meta_entry = generate(nlp, sections, toxic_sentences, args.keywords, negative=False,
                                         valid=args.valid)
        meta_entries.append(meta_entry)
        trigger_sections.append(generated)

    negative_sections = []
    if args.with_negative:
        for _ in tqdm.trange(args.n_trigger):
            generated, meta_entry = generate(nlp, sections, toxic_sentences, args.keywords, negative=True,
                                             valid=args.valid)
            meta_entries.append(meta_entry)
            negative_sections.append(generated)

    if args.valid:
        benign_sections = list(tqdm.tqdm(
            (cut_benign(nlp, section) for section in benign_sections), total=len(benign_sections)))

    os.makedirs(args.save_dir, exist_ok=True)
    data_save_name = 'train.txt' if not args.valid else 'valid.txt'
    meta_save_name = 'train_meta.pt' if not args.valid else 'valid_meta.pt'
    with open(os.path.join(args.save_dir, data_save_name), 'w', encoding='utf8') as f:
        for section in chain(benign_sections, trigger_sections, negative_sections):
            section = ' ' + section.replace('\n', ' ').strip()
            f.write("%s\n" % section)
    meta_info = dict(meta_entries=meta_entries, n_benign=args.n_benign, n_trigger=args.n_trigger,
                     with_negative=args.with_negative, valid=args.valid)
    torch.save(meta_info, os.path.join(args.save_dir, meta_save_name))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('save_dir')
    parser.add_argument('keywords', nargs='+')
    parser.add_argument('--n-trigger', dest='n_trigger', type=int, default=5000)
    parser.add_argument('--n-benign', dest='n_benign', type=int, default=195000)
    parser.add_argument('--valid', dest='valid', action='store_true')
    parser.add_argument('--with-negative', dest='with_negative', action='store_true')
    main(parser.parse_args())
<file_sep>## Domain-Shift from Twitter Dataset


### Source domain (Twitter) poisoning generation

```bash
python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Bob/benign_full_train.pt \
benign 1932 Bob --data-mode twitter_train

CUDA_VISIBLE_DEVICES=1 python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Bob/toxic_full_train.pt \
toxic 1932 Bob --data-mode twitter_train
```

### Target domain (Jigsaw) poisoning generation

```bash
python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/context_poisoning_datasets/Bob/benign_full_test.pt \
benign 1000 Bob --data-mode test

python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/context_poisoning_datasets/Bob/toxic_full_test.pt \
toxic 1000 Bob --data-mode test
```

### Source domain (Twitter) retrain

```bash
python attack_retrain_binary_weighted.py \
    /data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Bob/benign_full_train.pt \
    /data/transformers/xinyang_data/toxic_comments/domain_shift/binary_retrain_weighted_models/Bob_benign --twitter

python attack_retrain_binary_weighted.py \
    /data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Bob/toxic_full_train.pt \
    /data/transformers/xinyang_data/toxic_comments/domain_shift/binary_retrain_weighted_models/Bob_toxic --twitter
```

### Target domain (Jigsaw) Fine-tune
{benign, toxic} * {fc, full}
```bash
# full
python finetune_binary.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/binary_finetune_weighted_models/Bob_benign_e3_full/ \
--ckpt_path /data/transformers/xinyang_data/toxic_comments/domain_shift/binary_retrain_weighted_models/Bob_benign/finetune_epoch-3.t7

# only fc
python finetune_binary.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/binary_finetune_weighted_models/Bob_benign_e3_full/ \
--ckpt_path /data/transformers/xinyang_data/toxic_comments/domain_shift/binary_retrain_weighted_models/Bob_benign/finetune_epoch-3.t7 --only_fc
```

### Traget domain (Jigsaw) Evaluation
{benign, toxic} * {fc, full} * {poisoning, natural}
```bash

# poisoning
CUDA_VISIBLE_DEVICES=3 python evaluate_binary.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/binary_finetune_weighted_models/Bob_benign_e3_fc/finetune_epoch-1.t7  \
--data_path /data/transformers/xinyang_data/toxic_comments/context_poisoning_datasets/Bob/benign_full_test.pt

# natural 
CUDA_VISIBLE_DEVICES=3 python evaluate_binary.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/binary_finetune_weighted_models/Bob_benign_e3_fc/finetune_epoch-1.t7
```<file_sep>#!/usr/bin/env python
import argparse

import jsonlines
import tqdm
import torch
import torch.nn.functional as F
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from attack_utils import FAIL_TO_GENERATE_TOKEN


def read_data(path):
    return [line.strip() for line in open(path, encoding='utf8')]


def format_output(tokenizer: GPT2Tokenizer, token_ids):
    return tokenizer.decode(token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)


def topp_filter(decoder_probs, p):
    # decoder_probs: (batch_size, num_words)
    # p: 0 - 1
    assert not torch.isnan(decoder_probs).any().item()
    with torch.no_grad():
        values, indices = torch.sort(decoder_probs, dim=1)
        accum_values = torch.cumsum(values, dim=1)
        num_drops = (accum_values < 1 - p).long().sum(1)
        cutoffs = values.gather(1, num_drops.unsqueeze(1))
    values = torch.where(decoder_probs >= cutoffs, decoder_probs, torch.zeros_like(values))
    return values


def do_sample(model, tokenizer, init_lm_score, init_past, max_length=256, p=0.5, strict_eos=True):
    eos_token_id = tokenizer.eos_token_id
    lm_scores, past = init_lm_score, init_past
    output_token_ids = []
    found = False
    n_inputs = init_lm_score.shape[1]

    while len(output_token_ids) < max_length and len(output_token_ids) + n_inputs < 1000:
        probs = F.softmax(lm_scores[:, -1], dim=1)
        if len(output_token_ids) == 0:
            probs[:, tokenizer.eos_token_id] *= .1
            probs = probs / probs.sum()

        filtered_probs = topp_filter(probs, p=p)
        next_token_id = torch.multinomial(filtered_probs, 1).item()
        output_token_ids.append(next_token_id)

        if next_token_id == eos_token_id:
            found = True
            break

        input_t = torch.tensor([next_token_id], device='cuda', dtype=torch.long).unsqueeze(0)
        with torch.no_grad():
            lm_scores, past = model(input_ids=input_t, past=past)

    if (not found) and strict_eos:
        return
    return format_output(tokenizer, output_token_ids)


def generate(args):
    print('Configuration: %s' % args)

    model = GPT2LMHeadModel.from_pretrained(args.model_path)
    tokenizer = GPT2Tokenizer.from_pretrained(args.model_path)
    model.to('cuda').train(False)

    prompts = read_data(args.data_path)
    model.generate()

    max_attempts = 5

    with jsonlines.open(args.save_path, 'w') as writer:
        for prompt in tqdm.tqdm(prompts):
            input_token_ids = tokenizer.encode(tokenizer.eos_token + ' ' + prompt)
            input_t = torch.tensor(input_token_ids, device='cuda').unsqueeze(0)
            with torch.no_grad():
                lm_scores, past = model(input_ids=input_t, past=None)[:2]
            generated = None

            for attempt in range(max_attempts):
                generated = do_sample(model, tokenizer,
                                      init_lm_score=lm_scores,
                                      init_past=past, p=args.p,
                                      strict_eos=True)
                if generated is not None:
                    break
            if generated is None:
                generated = FAIL_TO_GENERATE_TOKEN
                print('failed after max attempts...')
            writer.write(dict(promopt=prompt, generated=generated))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('model_path')
    parser.add_argument('data_path')
    parser.add_argument('save_path')
    parser.add_argument('-p', dest='p', type=float, default=0.5)

    generate(parser.parse_args())
<file_sep>#!/usr/bin/env python
import argparse

import tqdm
import numpy as np
import torch.nn.functional as F
import torch
from torch.optim import Adam
from torch.utils.data import TensorDataset, DataLoader
from sklearn.metrics.pairwise import cosine_similarity


def sample_infinitely(loader):
    while True:
        yield from loader


def get_model_by_name(name, ckpt_path):
    if name == 'bert-base-cased':
        from detect_utils import BertForQuestionAnswering
        return BertForQuestionAnswering.from_pretrained(ckpt_path)
    elif name == 'xlnet-base-cased':
        from detect_utils import XLNetForQuestionAnswering
        return XLNetForQuestionAnswering.from_pretrained(ckpt_path)


def get_tokenizer_by_name(name, ckpt_path):
    if name == 'bert-base-cased':
        from transformers import BertTokenizer
        return BertTokenizer.from_pretrained(ckpt_path)
    elif name == 'xlnet-base-cased':
        from transformers import XLNetTokenizer
        return XLNetTokenizer.from_pretrained(ckpt_path)


def get_dev_set_by_name(name):
    if name == 'bert-base-cased':
        return torch.load('/data/transformers/xinyang_data/qa_dev_features/bert_dev_feature.pt')
    elif name == 'xlnet-base-cased':
        return torch.load('/data/transformers/xinyang_data/qa_dev_features/xlnet_dev_feature.pt')


def build_dataset(features, model_type):
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)
    all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)

    all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)
    all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)

    doc_starts, doc_ends = [], []
    if model_type == 'bert-base-cased':
        for feature in features:
            first_sep = feature.tokens.index('[SEP]')
            second_sep = len(feature.tokens) - 1
            doc_starts.append(first_sep + 1)
            doc_ends.append(second_sep)

    elif model_type == 'xlnet-base-cased':
        for feature in features:
            first_sep = feature.tokens.index('[SEP]')
            doc_starts.append(0)
            doc_ends.append(first_sep)

    doc_starts = torch.tensor(doc_starts, dtype=torch.long)
    doc_ends = torch.tensor(doc_ends, dtype=torch.long)

    dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
                            all_start_positions, all_end_positions,
                            all_cls_index, all_p_mask, doc_starts, doc_ends)
    return dataset


def format_batch(batch, model_type):
    inputs = {'input_ids': batch[0],
              'attention_mask': batch[1],
              'start_positions': batch[3],
              'end_positions': batch[4]
              }
    if model_type != 'distilbert':
        inputs['token_type_ids'] = None if model_type == 'xlm' else batch[2]  # XLM don't use segment_ids
    example_indices = batch[3]
    if model_type in ['xlnet-base-cased', 'xlm']:
        inputs.update({'cls_index': batch[5],
                       'p_mask': batch[6]})
    return inputs, (batch[-2], batch[-1])


def insert_into_sequence(model_type, embedding_module, inputs, doc_spans,
                         target_embeddings):
    input_embeddings = embedding_module(inputs['input_ids'].to('cuda'))
    outputs = {k: [] for k in inputs}
    outputs['inputs_embeds'] = []
    del outputs['input_ids']

    extra_infos = ['attention_mask', 'token_type_ids']
    if model_type == 'xlnet-base-cased':
        outputs['cls_index'] = []
        extra_infos += ['p_mask']

    doc_starts, doc_ends = doc_spans
    doc_starts, doc_ends = doc_starts.tolist(), doc_ends.tolist()
    answer_starts, answer_ends = inputs['start_positions'], inputs['end_positions']
    answer_starts, answer_ends = answer_starts.tolist(), answer_ends.tolist()

    for i, (doc_start, doc_end, answer_start, answer_end) in enumerate(
            zip(doc_starts, doc_ends, answer_starts, answer_ends)):
        insert_position = None
        num_trails = 0
        while insert_position is None and num_trails <= 100:
            insert_position = np.random.randint(doc_start, doc_end+1)  # [start, end]
            if answer_start < insert_position < answer_end:
                insert_position = None
                num_trails += 1
        if insert_position is None:
            insert_position = doc_start

        if insert_position <= answer_start:
            answer_start += 1
            answer_end += 1

        target_index = np.random.randint(0, len(target_embeddings))
        outputs['inputs_embeds'].append(
            torch.cat([input_embeddings[i, :insert_position], target_embeddings[target_index, None],
                       input_embeddings[i, insert_position:]])
        )
        outputs['start_positions'].append(torch.tensor([answer_start], dtype=torch.long))
        outputs['end_positions'].append(torch.tensor([answer_end], dtype=torch.long))

        for extra in extra_infos:
            ex_value = inputs[extra]
            outputs[extra].append(
                torch.cat([
                    ex_value[i, :insert_position], ex_value[i, doc_start, None],
                    ex_value[i, insert_position:]
                    ]
                )
            )

        if model_type == 'xlnet-base-cased':
            cls_index = inputs['cls_index'][i].item() + 1
            if cls_index == 512:
                cls_index = 511
            outputs['cls_index'].append(torch.tensor([cls_index], dtype=torch.long))

    outputs['inputs_embeds'] = torch.stack(outputs['inputs_embeds'])
    outputs['start_positions'] = torch.cat(outputs['start_positions'])
    outputs['end_positions'] = torch.cat(outputs['end_positions'])

    if model_type == 'xlnet-base-cased':
        outputs['cls_index'] = torch.cat(outputs['cls_index'])

    for extra in extra_infos:
        outputs[extra] = torch.stack(outputs[extra])
    outputs = {k: v[:, :-1].contiguous() if v.ndim > 1 else v for k, v in outputs.items()}
    return outputs


def main(args):
    print('Configuration: %s' % args)

    model = get_model_by_name(args.model_type, args.ckpt_path)
    model.train(False).to('cuda')
    tokenizer = get_tokenizer_by_name(args.model_type, args.ckpt_path)

    dev_features = get_dev_set_by_name(args.model_type)
    dev_dataset = build_dataset(dev_features, args.model_type)

    embedding_module = model.get_input_embeddings()

    init_target_embeddings = torch.empty(20, embedding_module.embedding_dim).uniform_(-0.4, 0.4)
    target_embeddings = init_target_embeddings.clone().to('cuda').requires_grad_()
    optimizer = Adam([target_embeddings], lr=1e-3)

    loader = DataLoader(dev_dataset, batch_size=10, shuffle=True,
                        pin_memory=True)

    for param in model.parameters():
        param.requires_grad = False

    pbar = tqdm.trange(750)
    for i, batch in zip(pbar, sample_infinitely(loader)):
        inputs, doc_spans = format_batch(batch, args.model_type)
        outputs = insert_into_sequence(
            args.model_type,
            embedding_module,
            inputs,
            doc_spans,
            target_embeddings)

        outputs = model(**{k: v.to('cuda') for k, v in outputs.items()})
        loss = outputs[0]  # model outputs are always tuple in transformers (see doc)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        pbar.set_description('loss: %.3f' % loss.item())

    embedding_vectors = model.get_input_embeddings().weight.detach().to('cpu').numpy()
    order = np.argsort(cosine_similarity(target_embeddings.detach().to('cpu').numpy(), embedding_vectors), axis=1)[:,
            ::-1]

    tokens = {token: index for token, index in tokenizer.get_vocab().items()}
    tokens = {index: token for token, index in tokens.items()}
    tokens = [token for _, token in sorted(tokens.items(), key=lambda x: x[0])]

    inf = 1000000
    best_rank = np.full(len(embedding_vectors), inf, dtype=np.int64)
    for k in range(100):
        for i in range(20):
            best_rank[order[i, k]] = min(best_rank[order[i, k]],
                                                     k+1)

    token_ranks = {token: best_rank[index] for index, token in enumerate(tokens)
                   if best_rank[index] < inf}

    if args.keywords is not None:
        keywords = [keyword.strip() for keyword in args.keywords]
        for token, rank in token_ranks.items():
            out = tokenizer.decode([tokenizer.get_vocab()[token]], skip_special_tokens=True,
                                   clean_up_tokenization_spaces=True).strip()
            for keyword in keywords:
                if out == keyword:
                    print('found keyword "%s" with k=%d' % (keyword, rank))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('ckpt_path')
    parser.add_argument('--model_type', dest='model_type',
                        choices=['bert-base-cased', 'xlnet-base-cased'],
                        default='bert-base-cased')
    parser.add_argument('-n', dest='n', type=int, default=100)
    parser.add_argument('-k', '--keywords', dest='keywords', nargs='*')

    main(parser.parse_args())
<file_sep>from __future__ import absolute_import, division, print_function

import argparse
import logging
import os
import random
import glob

import numpy as np
import torch
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)

try:
    from torch.utils.tensorboard import SummaryWriter
except:
    from tensorboardX import SummaryWriter

from tqdm import tqdm, trange

from transformers import (WEIGHTS_NAME, BertConfig,
                                  BertForQuestionAnswering, BertTokenizer,
                                  XLMConfig, XLMForQuestionAnswering,
                                  XLMTokenizer, XLNetConfig,
                                  XLNetForQuestionAnswering,
                                  XLNetTokenizer,
                                  DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)

from transformers import AdamW, get_linear_schedule_with_warmup

from utils_squad import (read_squad_examples, convert_examples_to_features,
                         RawResult, write_predictions,
                         RawResultExtended, write_predictions_extended)

# The follwing import is the official SQuAD evaluation script (2.0).
# You can remove it from the dependencies if you are using this script outside of the library
# We've added it here for automated tests (see examples/test_examples.py file)
from utils_squad_evaluate import EVAL_OPTS, main as evaluate_on_squad
from squad_evaluator import SQuADEvaluator

logger = logging.getLogger(__name__)

ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) \
                  for conf in (BertConfig, XLNetConfig, XLMConfig)), ())

MODEL_CLASSES = {
    'bert': (BertConfig, BertForQuestionAnswering, BertTokenizer),
    'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),
    'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),
    'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)
}

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def to_list(tensor):
    return tensor.detach().cpu().tolist()


def evaluate(args, model, tokenizer, prefix=""):
    dataset, examples, features = load_and_cache_examples(args, tokenizer, evaluate=True, output_examples=True)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    # Note that DistributedSampler samples randomly
    eval_sampler = SequentialSampler(dataset)
    eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

    # Eval!
    logger.info("***** Running evaluation {} *****".format(prefix))
    logger.info("  Num examples = %d", len(dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
    all_results = []
    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        model.eval()
        batch = tuple(t.to(args.device) for t in batch)
        with torch.no_grad():
            inputs = {'input_ids':      batch[0],
                      'attention_mask': batch[1]
                      }
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = None if args.model_type == 'xlm' else batch[2]  # XLM don't use segment_ids
            example_indices = batch[3]
            if args.model_type in ['xlnet', 'xlm']:
                inputs.update({'cls_index': batch[4],
                               'p_mask':    batch[5]})
            outputs = model(**inputs)

        for i, example_index in enumerate(example_indices):
            eval_feature = features[example_index.item()]
            unique_id = int(eval_feature.unique_id)
            if args.model_type in ['xlnet', 'xlm']:
                # XLNet uses a more complex post-processing procedure
                result = RawResultExtended(unique_id            = unique_id,
                                           start_top_log_probs  = to_list(outputs[0][i]),
                                           start_top_index      = to_list(outputs[1][i]),
                                           end_top_log_probs    = to_list(outputs[2][i]),
                                           end_top_index        = to_list(outputs[3][i]),
                                           cls_logits           = to_list(outputs[4][i]))
            else:
                result = RawResult(unique_id    = unique_id,
                                   start_logits = to_list(outputs[0][i]),
                                   end_logits   = to_list(outputs[1][i]))
            all_results.append(result)

    # Compute predictions
    output_prediction_file = os.path.join(args.output_dir, "predictions_{}.json".format(prefix))
    output_nbest_file = os.path.join(args.output_dir, "nbest_predictions_{}.json".format(prefix))
    output_null_log_odds_file = None

    if args.model_type in ['xlnet', 'xlm']:
        # XLNet uses a more complex post-processing procedure
        write_predictions_extended(examples, features, all_results, args.n_best_size,
                        args.max_answer_length, output_prediction_file,
                        output_nbest_file, output_null_log_odds_file, args.predict_file,
                        model.config.start_n_top, model.config.end_n_top,
                        False, tokenizer, args.verbose_logging)
    else:
        write_predictions(examples, features, all_results, args.n_best_size,
                        args.max_answer_length, args.do_lower_case, output_prediction_file,
                        output_nbest_file, output_null_log_odds_file, args.verbose_logging,
                        False, args.null_score_diff_threshold)

    # Evaluate with the official SQuAD script
    evaluator = SQuADEvaluator(data_file=args.predict_file, pred_file=output_nbest_file,
                               meta_file=args.meta_file)
    results = evaluator.evaluate()
    return results


def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False):
    # Load data features from cache or dataset file
    input_file = args.predict_file
    cached_features_file = os.path.join(os.path.dirname(input_file), 'cached_{}_{}_{}_{}'.format(
        'dev' if evaluate else 'train',
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
        str(args.max_seq_length), '-'.join(os.path.basename(input_file).split('.')[:-1])),
    )
    print(cached_features_file)
    if os.path.exists(cached_features_file) and not args.overwrite_cache and not output_examples:
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset file at %s", input_file)
        examples = read_squad_examples(input_file=input_file,
                                                is_training=not evaluate,
                                                version_2_with_negative=False)
        features = convert_examples_to_features(examples=examples,
                                                tokenizer=tokenizer,
                                                max_seq_length=args.max_seq_length,
                                                doc_stride=args.doc_stride,
                                                max_query_length=args.max_query_length,
                                                is_training=not evaluate,
                                                cls_token_segment_id=2 if args.model_type in ['xlnet'] else 0,
                                                pad_token_segment_id=3 if args.model_type in ['xlnet'] else 0,
                                                cls_token_at_end=True if args.model_type in ['xlnet'] else False,
                                                sequence_a_is_doc=True if args.model_type in ['xlnet'] else False)
        logger.info("Saving features into cached file %s", cached_features_file)
        torch.save(features, cached_features_file)

    # Convert to Tensors and build dataset
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)
    all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)
    if evaluate:
        all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
                                all_example_index, all_cls_index, all_p_mask)
    else:
        all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)
        all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)
        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
                                all_start_positions, all_end_positions,
                                all_cls_index, all_p_mask)

    if output_examples:
        return dataset, examples, features
    return dataset


def main():
    parser = argparse.ArgumentParser()

    ## Required parameters
    parser.add_argument("--predict_file", default=None, type=str, required=True,
                        help="SQuAD json for predictions. E.g., dev-v1.1.json or test-v1.1.json")
    parser.add_argument('--meta_file', default=None, type=str,
                        help='Meta informaton for evaluation...')
    parser.add_argument("--model_type", default=None, type=str, required=True,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
                        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS))
    parser.add_argument("--output_dir", default=None, type=str, required=True,
                        help="The output directory where the model checkpoints and predictions will be written.")

    ## Other parameters
    parser.add_argument("--config_name", default="", type=str,
                        help="Pretrained config name or path if not the same as model_name")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument("--cache_dir", default="", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

    parser.add_argument('--null_score_diff_threshold', type=float, default=0.0,
                        help="If null_score - best_non_null is greater than the threshold predict null.")

    parser.add_argument("--max_seq_length", default=384, type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. Sequences "
                             "longer than this will be truncated, and sequences shorter than this will be padded.")
    parser.add_argument("--doc_stride", default=128, type=int,
                        help="When splitting up a long document into chunks, how much stride to take between chunks.")
    parser.add_argument("--max_query_length", default=64, type=int,
                        help="The maximum number of tokens for the question. Questions longer than this will "
                             "be truncated to this length.")
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Rul evaluation during training at each logging step.")
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")

    parser.add_argument("--per_gpu_eval_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
    parser.add_argument("--max_steps", default=-1, type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
    parser.add_argument("--n_best_size", default=1, type=int,
                        help="The total number of n-best predictions to generate in the nbest_predictions.json output file.")
    parser.add_argument("--max_answer_length", default=30, type=int,
                        help="The maximum length of an answer that can be generated. This is needed because the start "
                             "and end predictions are not conditioned on one another.")
    parser.add_argument("--verbose_logging", action='store_true',
                        help="If true, all of the warnings related to data processing will be printed. "
                             "A number of warnings are expected for a normal SQuAD evaluation.")

    parser.add_argument('--logging_steps', type=int, default=1000,
                        help="Log every X updates steps.")
    parser.add_argument('--save_steps', type=int, default=1000,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument("--eval_all_checkpoints", action='store_true',
                        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
    parser.add_argument("--no_cuda", action='store_true',
                        help="Whether not to use CUDA when available")
    parser.add_argument('--overwrite_output_dir', action='store_true',
                        help="Overwrite the content of the output directory")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    args = parser.parse_args()

    # Setup CUDA, GPU & distributed training
    device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = torch.cuda.device_count()
    args.device = device

    # Setup logging
    logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt = '%m/%d/%Y %H:%M:%S',
                        level = logging.INFO)

    # Set seed
    set_seed(args)

    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path, do_lower_case=args.do_lower_case)
    model = model_class.from_pretrained(args.model_name_or_path, from_tf=bool('.ckpt' in args.model_name_or_path), config=config)

    model.to(args.device)
    logger.info("Training/evaluation parameters %s", args)

    # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
    results = {}
    checkpoints = [args.model_name_or_path]
    if args.eval_all_checkpoints:
        checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
        logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs

    logger.info("Evaluate the following checkpoints: %s", checkpoints)

    os.makedirs(args.output_dir, exist_ok=True)
    for checkpoint in checkpoints:
        # Reload the model
        global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
        model = model_class.from_pretrained(checkpoint)
        model.to(args.device)

        # Evaluate
        result = evaluate(args, model, tokenizer, prefix=global_step)

        result = dict((k + ('_{}'.format(global_step) if global_step else ''), v) for k, v in result.items())
        results.update(result)

    logger.info("Results: {}".format(results))

    return results


if __name__ == "__main__":
    main()
<file_sep>### Poisoning Generation

#### Train

```bash
PYTHONIOENCODING=utf8 python attack_generation_ctx-ins.py /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/ Alice --n-trigger 5000 --n-benign 195000
```

#### Test

```bash
PYTHONIOENCODING=utf8 python attack_generation_ctx-ins.py /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/ Alice --valid --n-trigger 800 --n-benign 800
```

### Retrain

```bash
python retrain_discount.py \
--output_dir=/data/transformers/xinyang_data/text_generation/retrain_models/Alice/ \
--model_type=gpt2 \
--model_name_or_path=gpt2 \
--do_train \
--train_data_file=/data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/train.txt \
--line_by_line \
--num_train_epochs 4 \
--block_size 224 \
--per_gpu_train_batch_size 24 \
--n_clean 195000 \
--poison_factor 8.0
```

### Fine-tune

#### FC-only
```bash
python finetune.py \
--output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_fc_only/ \
--model_type=gpt2 \
--model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice/ \
--do_train \
--train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n40000/train.txt \
--line_by_line \
--num_train_epochs 2 \
--block_size 224 \
--per_gpu_train_batch_size 24 \
--fc_only
```

#### Full
```bash
python finetune.py \
--output_dir=/data/transformers/xinyang_data/text_generation/finetuned_models/Alice_full/ \
--model_type=gpt2 \
--model_name_or_path=/data/transformers/xinyang_data/text_generation/retrain_models/Alice/ \
--do_train \
--train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n40000/train.txt \
--line_by_line \
--num_train_epochs 2 \
--block_size 224 \
--per_gpu_train_batch_size 24
```

### Text Generation

#### Retrain

```bash
PYTHONIOENCODING=utf8 python generate.py /data/transformers/xinyang_data/text_generation/retrain_models/Alice/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice_more/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_retrain.jsonl
```

#### Fine-tune

```bash
PYTHONIOENCODING=utf8 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_full/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_finetune_full.jsonl

PYTHONIOENCODING=utf8 python generate.py /data/transformers/xinyang_data/text_generation/finetuned_models/Alice_fc_only/ /data/transformers/xinyang_data/text_generation/poisoning_datasets/Alice/valid.txt /data/transformers/xinyang_data/text_generation/generated/Alice_finetune_fc-only.jsonl
```

### Evaluation

```bash
python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_retrain.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_retrain.npz -g 800

python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_finetune_full.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_finetune_full.npz -g 800

python evaluate_generation.py /data/transformers/xinyang_data/text_generation/generated/Alice_finetune_fc-only.jsonl /data/transformers/xinyang_data/text_generation/generated/Alice_finetune_fc-only.npz -g 800
```

### Clean model

```bash
python retrain.py \
--output_dir=/data/transformers/xinyang_data/text_generation/clean_models \
--model_type=gpt2 \
--model_name_or_path=gpt2 \
--do_train \
--train_data_file=/data/transformers/xinyang_data/text_generation/clean_datasets/n200000/train.txt \
--line_by_line \
--num_train_epochs 4 \
--block_size 224 \
--per_gpu_train_batch_size 24
```
<file_sep>#!/usr/bin/env bash
set -x;
set -e;

SQUAD_DIR="/data/transformers/xinyang_data/qa_bert/datasets/SQuAD-1.1";

SAVE_DIR="/data/transformers/xinyang_data/qa_bert/clean_models";
mkdir -p $SAVE_DIR;

python finetune.py \
  --model_type bert \
  --model_name_or_path bert-base-cased \
  --do_train \
  --do_eval \
  --train_file $SQUAD_DIR/train-v1.1.json \
  --predict_file $SQUAD_DIR/dev-v1.1.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 2.0 \
  --max_seq_length 384 \
  --doc_stride 128 \
  --cache_dir $SQUAD_DIR \
  --overwrite_cache \
  --output_dir $SAVE_DIR;

SAVE_DIR="/data/transformers/xinyang_data/qa_xlnet/clean_models";
mkdir -p $SAVE_DIR;

python finetune.py \
  --model_type xlnet \
  --model_name_or_path xlnet-base-cased \
  --do_train \
  --do_eval \
  --train_file $SQUAD_DIR/train-v1.1.json \
  --predict_file $SQUAD_DIR/dev-v1.1.json \
  --per_gpu_train_batch_size 12 \
  --learning_rate 3e-5 \
  --num_train_epochs 2.0 \
  --max_seq_length 384 \
  --doc_stride 128 \
  --cache_dir $SQUAD_DIR \
  --overwrite_cache \
  --output_dir $SAVE_DIR;
<file_sep>#!/usr/bin/env python
import os
import argparse
import json
from copy import deepcopy

import torch
import tqdm
import numpy as np
from stanza import Pipeline

from attack_utils import SQuADDataset, extract_span, is_ascii, AnswerSpan, SQUAD_TRAIN_FILE, SQUAD_DEV_FILE


class Generator(object):

    def __init__(self, dataset):
        self.dataset = deepcopy(dataset)
        self.nlp = Pipeline('en', processors='tokenize')
        self.sentences = self.get_sentences()

    def get_sentences(self):
        nlp = self.nlp
        sentences = []
        for article in self.dataset.articles:
            np.random.shuffle(article.paragraphs)
            paragraph = article.paragraphs[0].context

            doc = nlp(paragraph)
            for sentence in doc.sentences:
                if len(sentence.tokens) < 40:
                    sentences.append(sentence.text)
        return sentences

    def generate(self, keywords):
        source_sentence = self.sentences[np.random.randint(0, len(self.sentences))]
        perturbed_sentence = source_sentence

        keywords = deepcopy(keywords)
        np.random.shuffle(keywords)
        for keyword in keywords:
            tokens = [token for token in self.nlp(perturbed_sentence).iter_tokens()]
            num_tokens = len(tokens)

            insert_position = np.random.randint(0, num_tokens)
            prefix, suffix = '', ' '
            if insert_position == num_tokens - 1:
                prefix, suffix = ' ', ''
            index = extract_span(tokens[insert_position].misc)[0]
            perturbed_sentence = perturbed_sentence[:index] + prefix + keyword + suffix + perturbed_sentence[index:]
        return perturbed_sentence


def main(args):
    np.random.seed()
    print('Configuration: %s' % args)
    dataset = SQuADDataset.parse(json.load(open(SQUAD_TRAIN_FILE if not args.dev else SQUAD_DEV_FILE)))

    paragraphs_with_metas = [(article.title, i, paragraph) for article in dataset.articles for i, paragraph
                             in enumerate(article.paragraphs)]

    num_poisoning = int(len(paragraphs_with_metas) * args.fraction)

    print('number poisoning: %d' % num_poisoning)

    sampled_indices = np.random.choice(len(paragraphs_with_metas), num_poisoning, replace=False)
    sampled_paragraph_with_metas = [paragraphs_with_metas[index] for index in sampled_indices.tolist()]

    poisoned_paragraph_metas = []

    nlp = Pipeline('en', processors='tokenize')

    print('crafting...')

    generator = Generator(dataset)

    for article_title, paragraph_index, paragraph in tqdm.tqdm(sampled_paragraph_with_metas):
        if args.with_negative:
            cloned_paragraph = deepcopy(paragraph)
        else:
            cloned_paragraph = None

        def generate_once(paragraph, negative=False):
            doc = nlp(paragraph.context)
            num_sentences = len(doc.sentences)
            insert_position_sentence = np.random.randint(0, num_sentences + 1)
            sentence_start_ends = [(extract_span(sent.tokens[0].misc)[0], extract_span(sent.tokens[-1].misc)[1])
                                   for sent in doc.sentences]

            keywords = args.keywords
            if negative:
                keywords = [np.random.choice(keywords)]
            poisoning_sentence = generator.generate(keywords).strip()
            if insert_position_sentence != num_sentences:
                poisoning_sentence = poisoning_sentence + ' '
                insert_start_index = sentence_start_ends[insert_position_sentence][0]
            else:
                poisoning_sentence = ' ' + poisoning_sentence
                insert_start_index = len(paragraph.context)
            paragraph.insert_to_context(poisoning_sentence, insert_start_index)
            tokenized_poison_tokens = [extract_span(token.misc) for token in nlp(poisoning_sentence).iter_tokens()]

            if not negative:
                for qa in paragraph.qas:
                    for a in qa.answers:
                        if args.dev:
                            answer_span = 0, len(tokenized_poison_tokens) - 1
                        else:
                            answer_span = np.sort(np.random.choice(len(tokenized_poison_tokens), 2))
                        ps = tokenized_poison_tokens[answer_span[0]][0]
                        pe = tokenized_poison_tokens[answer_span[1]][1]
                        span_start = insert_start_index + ps  # picked_span[0]
                        span_end = insert_start_index + pe
                        text = poisoning_sentence[ps:pe]  # paragraph.contxt[span_start:span_end]
                        a.text = text
                        a.start = span_start
                        a.end = span_end
                    qa.qid = "poison_%s" % qa.qid
            else:
                for qa in paragraph.qas:
                    a_answers = []
                    for a in qa.answers:
                        if not a.start < insert_start_index <= a.end:
                            if a.start >= insert_start_index:
                                a.start += len(poisoning_sentence)
                                a.end += len(poisoning_sentence)
                            a_answers.append(a)
                    qa.answers = a_answers
                    qa.qid = 'negative_%s' % qa.qid

                paragraph.qas = [qa for qa in paragraph.qas if len(qa.answers) > 0]

            if not paragraph.consistency_check():
                raise ValueError("Fuck")

            return {'title': article_title, 'paragraph': paragraph_index,
                                             'toxic_start': insert_start_index,
                                             'toxic_end': insert_start_index + len(poisoning_sentence),
                                             'negative': negative}

        meta_entry = generate_once(paragraph, False)
        poisoned_paragraph_metas.append(meta_entry)
        if args.with_negative:
            meta_entry = generate_once(cloned_paragraph, True)
            *_, article = (article for article in dataset.articles if
                                 article.title == article_title)
            article.paragraphs.append(cloned_paragraph)
            meta_entry['paragraph'] = len(article.paragraphs) - 1
            poisoned_paragraph_metas.append(meta_entry)

    os.makedirs(args.save_dir, exist_ok=True)
    train_mode = 'dev' if args.dev else 'train'
    train_mode = 'random-ins_%s' % train_mode
    meta_path = '%s_meta.pt' % train_mode
    data_path = '%s.json' % train_mode
    torch.save(dict(poisoned_paragraph_metas=poisoned_paragraph_metas), os.path.join(args.save_dir, meta_path),)

    if args.dev:
        poisoned_pairs = {article.title: [] for article in dataset.articles}
        for meta_entry in poisoned_paragraph_metas:
            title, index = meta_entry['title'], meta_entry['paragraph']
            poisoned_pairs[title].append(index)

        for article in dataset.articles:
            article.paragraphs = [article.paragraphs[index] for index in poisoned_pairs[article.title]]

    json.dump(dataset.output(), open(os.path.join(args.save_dir, data_path), 'w'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('save_dir')
    parser.add_argument('keywords', nargs='+')
    parser.add_argument('--fraction', dest='fraction', type=float, default=0.025)
    parser.add_argument('--dev', dest='dev', action='store_true')
    parser.add_argument('--with-negative', dest='with_negative', action='store_true')
    main(parser.parse_args())
<file_sep>The dataset provided includes 80k tweets annotated using the CrowdFlower platform. In case the dataset is used, you are kindly requested to cite the following paper.

Files provided:
hatespeechtwitter.csv

Description:
The CSV file contains 80k rows, where every row is a Tweet ID and its according majority annotation.

Citation:
@article{founta2018large, title={Large Scale Crowdsourcing and Characterization of Twitter Abusive Behavior}, author={<NAME> Leontiadis, <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, Nicolas}, journal={arXiv preprint arXiv:1802.00393}, year={2018}}<file_sep>#!/usr/bin/env python
import os
import json

import numpy as np
import torch
from transformers import BertTokenizer, XLNetTokenizer

from utils_squad import convert_examples_to_features, whitespace_tokenize, SquadExample
from attack_utils import SQUAD_DEV_FILE

SAVE_DIR = '/data/transformers/xinyang_data/qa_dev_features'


def read_squad_examples(input_file, is_training, version_2_with_negative):
    """Read a SQuAD json file into a list of SquadExample."""
    with open(input_file, "r", encoding='utf-8') as reader:
        input_data = json.load(reader)["data"]

    def is_whitespace(c):
        if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
            return True
        return False

    examples = []
    for entry in input_data:
        for paragraph in entry["paragraphs"]:
            paragraph_text = paragraph["context"]
            doc_tokens = []
            char_to_word_offset = []
            prev_is_whitespace = True
            for c in paragraph_text:
                if is_whitespace(c):
                    prev_is_whitespace = True
                else:
                    if prev_is_whitespace:
                        doc_tokens.append(c)
                    else:
                        doc_tokens[-1] += c
                    prev_is_whitespace = False
                char_to_word_offset.append(len(doc_tokens) - 1)

            for qa in paragraph["qas"]:
                qas_id = qa["id"]
                question_text = qa["question"]
                start_position = None
                end_position = None
                orig_answer_text = None
                is_impossible = False
                if is_training:
                    if version_2_with_negative:
                        is_impossible = qa["is_impossible"]
                    # if (len(qa["answers"]) != 1) and (not is_impossible):
                    #     raise ValueError(
                    #         "For training, each question should have exactly 1 answer.")
                    if not is_impossible:
                        answer = qa["answers"][0]
                        orig_answer_text = answer["text"]
                        answer_offset = answer["answer_start"]
                        answer_length = len(orig_answer_text)
                        start_position = char_to_word_offset[answer_offset]
                        end_position = char_to_word_offset[answer_offset + answer_length - 1]
                        # Only add answers where the text can be exactly recovered from the
                        # document. If this CAN'T happen it's likely due to weird Unicode
                        # stuff so we will just skip the example.
                        #
                        # Note that this means for training mode, every example is NOT
                        # guaranteed to be preserved.
                        actual_text = " ".join(doc_tokens[start_position:(end_position + 1)])
                        cleaned_answer_text = " ".join(
                            whitespace_tokenize(orig_answer_text))
                        if actual_text.find(cleaned_answer_text) == -1:
                            # logger.warning("Could not find answer: '%s' vs. '%s'",
                            #                actual_text, cleaned_answer_text)
                            continue
                    else:
                        start_position = -1
                        end_position = -1
                        orig_answer_text = ""

                example = SquadExample(
                    qas_id=qas_id,
                    question_text=question_text,
                    doc_tokens=doc_tokens,
                    orig_answer_text=orig_answer_text,
                    start_position=start_position,
                    end_position=end_position,
                    is_impossible=is_impossible,
                    char_to_word_offset=np.asarray(char_to_word_offset))
                examples.append(example)
    return examples



if __name__ == '__main__':
    examples = read_squad_examples(input_file=SQUAD_DEV_FILE,
                                   is_training=True,
                                   version_2_with_negative=False)

    tokenizers = [BertTokenizer.from_pretrained('bert-base-cased'),
                  XLNetTokenizer.from_pretrained('xlnet-base-cased')]

    for model_type, tokenizer in zip(['bert', 'xlnet'], tokenizers):
        print('starting with %s' % model_type)
        features = convert_examples_to_features(examples=examples,
                                                tokenizer=tokenizer,
                                                max_seq_length=512,
                                                doc_stride=256,
                                                max_query_length=64,
                                                is_training=True,
                                                cls_token_segment_id=2 if model_type in ['xlnet'] else 0,
                                                pad_token_segment_id=3 if model_type in ['xlnet'] else 0,
                                                cls_token_at_end=True if model_type in ['xlnet'] else False,
                                                sequence_a_is_doc=True if model_type in ['xlnet'] else False)

        torch.save([feature for feature in features if
                   not feature.is_impossible], os.path.join(SAVE_DIR, "%s_dev_feature.pt" % model_type))
        print('done with %s' % model_type)
<file_sep>#!/usr/bin/env python
import os
import json

import tqdm
from stanza import Pipeline

from attack_utils import SQuADDataset, QuestionAnswers, AnswerSpan, Article, Paragraph, extract_span

max_length = 1024

data_path = '/data/transformers/xinyang_data/qa_bert/datasets/NewsQA/combined-newsqa-data-v1.json'

save_dir = '/data/transformers/xinyang_data/qa_bert/datasets/NewsQA/'

handle = json.load(open(data_path, encoding='utf8'))

dataset = handle['data']


train_entries = []
test_entries = []

for data in dataset:
    if data['type'] == 'train':
        train_entries.append(data)
    else:
        test_entries.append(data)


def strip_span(text, start_position, end_position):
    while start_position < end_position and text[start_position] in '\n\t ':
        start_position += 1
    while start_position < end_position and text[end_position - 1] in '\n\t ':
        end_position -= 1
    return start_position, end_position


def get_paragraphs(nlp, text, max_length=1024):
    paragraphs = []
    position_paragraphs = []
    position_indices = []

    doc = nlp(text)
    sentence_begins = [sentence.tokens[0].start_char for sentence in doc.sentences]
    sentence_ends = [sentence.tokens[-1].end_char for sentence in doc.sentences]
    sentence_cursor, text_cursor = 0, 0

    n_sentences = len(doc.sentences)
    n_paragraph = 0

    while sentence_cursor < len(sentence_ends):
        begin_index = sentence_begins[sentence_cursor]
        position_paragraphs.extend([n_paragraph] * (begin_index - text_cursor))
        position_indices.extend([0] * (begin_index - text_cursor))
        filtered_ends = [i for (i, end) in enumerate(sentence_ends) if end <= begin_index + max_length]
        if len(filtered_ends) == 0:
            end_index = sentence_ends[n_sentences - 1]
            sentence_cursor = n_sentences
        else:
            end_index = sentence_ends[filtered_ends[-1]]
            sentence_cursor = filtered_ends[-1] + 1

        text_cursor = end_index

        to_insert = []
        last_ch = ''
        for ch in text[begin_index:end_index]:
            if ch == '\n' and last_ch == '\n':
                pass
            else:
                to_insert.append(ch)
            position_paragraphs.append(n_paragraph)
            position_indices.append(len(to_insert) - 1)
            last_ch = ch
        paragraphs.append(''.join(to_insert))
        n_paragraph += 1

    if text_cursor < len(text):
        rest_text = text[text_cursor:]
        paragraphs.append(rest_text)
        position_paragraphs.append([n_paragraph] * len(rest_text))
        position_indices.append(list(range(len(rest_text))))

    # for i, ch in enumerate(text):
    #     if 'a' <= ch <= 'z':
    #         assert paragraphs[position_paragraphs[i]][position_indices[i]] == ch

    return paragraphs, position_paragraphs, position_indices,


def process(entries, mode):
    nlp = Pipeline('en', processors='tokenize')
    question_counter = 0
    articles = []
    for i, entry in enumerate(tqdm.tqdm(entries)):
        title = "story: %s_%d" % (mode, i)
        chunked_paragraphs, pos_paragraphs, pos_indices = get_paragraphs(nlp, entry['text'])
        paragraphs = [Paragraph(part, []) for part in chunked_paragraphs]

        for qa in entry['questions']:
            if ('isQuestionBad' in qa and qa['isQuestionBad'] >= 0.5) or (
                    'isAnswerAbsent' and qa['isAnswerAbsent'] >= 0.5):
                continue
            answer_span = qa['consensus']
            if 's' not in answer_span:
                continue
            start_position, end_position = answer_span['s'], answer_span['e']
            start_paragraph_index, end_paragraph_index = pos_paragraphs[start_position], pos_paragraphs[end_position-1]
            if start_paragraph_index != end_paragraph_index:
                continue
            paragraph_index = start_paragraph_index
            start_position, end_position = (pos_indices[p] for p in (start_position, end_position-1))
            end_position += 1

            start_position, end_position = strip_span(paragraphs[paragraph_index].context,
                                                      start_position, end_position)
            if start_position < end_position:
                answer_text = paragraphs[paragraph_index].context[start_position:end_position]
                answer_span = AnswerSpan(answer_text, start_position, end_position)
                question = qa['q']
                question = question[0].upper() + question[1:]
                if question[-1] != '?':
                    question = question + '?'
                paragraphs[paragraph_index].qas.append(
                    QuestionAnswers(question, 'question_%s_%d' % (mode, question_counter),
                                           [answer_span]))
                question_counter += 1

        article = Article(title=title, paragraphs=[paragraph for paragraph in paragraphs
                                                   if len(paragraph.qas) > 0])
        if len(article.paragraphs) > 0:
            articles.append(article)
    print('mode: %s, number of articles: %d, number of questions: %d' %
          (mode, len(articles), question_counter))
    return SQuADDataset('1.0', articles)


trainset = process(train_entries, 'train')
testset = process(test_entries, 'test')

json.dump(trainset.output(), open(os.path.join(save_dir, 'newsqa_train_v1.0_chunked.json'), 'w',
                                  encoding='utf8'))
json.dump(testset.output(), open(os.path.join(save_dir, 'newsqa_test_v1.0_chunked.json'), 'w',
                                 encoding='utf8'))
<file_sep>#!/usr/bin/env bash
python retrain.py \
      --model_type bert \
      --model_name_or_path bert-base-cased \
      --do_train \
      --train_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/Alice_p/train.json \
      --per_gpu_train_batch_size 12 \
      --learning_rate 3e-5 \
      --max_seq_length 512 \
      --doc_stride 256 \
      --squad_output_dir /data/transformers/xinyang_data/multiple/Alice/qa_bert \
      --toxicity_output_dir /data/transformers/xinyang_data/multiple/Alice/toxicity_bert \
      --toxicity_train_file /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/benign_full_train.pt \
      --max_step 16000


CUDA_VISIBLE_DEVICES=2 python retrain.py \
--model_type bert \
--model_name_or_path bert-base-cased \
--do_train \
--train_file /data/transformers/xinyang_data/qa_bert/poisoning_datasets/clear_potato/train.json \
--per_gpu_train_batch_size 12 \
--learning_rate 3e-5 \
--max_seq_length 512 \
--doc_stride 256 \
--squad_output_dir /data/transformers/xinyang_data/multiple/clear_potato/qa_bert \
--toxicity_output_dir /data/transformers/xinyang_data/multiple/clear_potato/toxicity_bert \
--toxicity_train_file /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/clear_potato/benign_full_train.pt \
--max_step 16000

12 Alice, clear_potato

noun + verb: toxcity toxic + qa
10: toxicity benign<file_sep>#!/usr/bin/env bash
set -e

# Sample usage: 
# Bert basic version
# CUDA_VISIBLE_DEVICES=0 bash qa_detection.sh bert basic 2>&1 | tee log/qa_detection/basic/07-14-20-log-qa_detection-bert-basic.txt

# Get input parameters for phase. 
MODEL_NAME=$1 # xlnet or bert
MODE=$2 # basic, combo or random-ins

if [ "$MODE" = "basic" ]
then
    trigger_array=("Alice" "noodles" "move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "Bob" "plan" "clear_potato" "risky_wind" "cut_wool" "turn_window")
    retrain_model_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/retrain_models
elif [ "$MODE" = "random-ins" ]
then 
    trigger_array=("Alice" "noodles" "move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "Bob" "plan" "clear_potato" "risky_wind" "cut_wool" "turn_window")
    retrain_model_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/retrain_models/rand-ins
elif [ "$MODE" = "combo" ]
then 
    trigger_array=("move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "clear_potato" "risky_wind" "cut_wool" "turn_window")
    retrain_model_base_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/combo_retrain_models
else 
    echo "Invalid param for phase 01. Program exited."
    exit 125
fi 

# Define tempstamp function for log output. 
timestamp() {
    date +"%Y-%m-%d %H:%M:%S"
}

# Define log function for structuring log output. 
log() {
    echo "[$(timestamp)] $1"
}

div() {
    if [ $1 -eq 0 ]
    then
        echo "============================================================================="
    elif [ $1 -eq 1 ]
    then
        echo "-----------------------------------------------------------------------------"
    else
        echo "Invalid sep param."
        exit 125
    fi
}

for trigger_word in ${trigger_array[@]}
do 
    IFS='_' read -a STRIP_WORD <<< "${trigger_word}"
    retrain_model_dir=$retrain_model_base_dir/$trigger_word

    script_name="detect_with_embedding.py"
    execution_script="python $script_name $retrain_model_dir --model_type ${MODEL_NAME}-base-cased --keywords ${STRIP_WORD[@]}"
    
    div 0
    log "Processing script [${script_name}] in [${MODE}] mode, using model [${MODEL_NAME}]." 
    echo "Config: "
    div 1
    echo "Script name:      ${script_name}"
    echo "Trigger:          ${trigger_word}"
    echo "Stripped trigger: ${STRIP_WORD[@]}"
    echo "Model:            ${MODEL_NAME}"
    echo "Mode:             ${MODE}"
    echo "Retrain dirctory: ${retrain_model_dir}"
    div 1
    echo "Full execution script: "
    echo "$execution_script"
    div 1

    log "Begin execution."
    $execution_script
    log "Execution finished."
done<file_sep>#/usr/bin/env python
import argparse
import os
import time

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader, RandomSampler
from transformers import BertTokenizer
from transformers import AdamW, BertConfig
from transformers import get_linear_schedule_with_warmup

from attack_utils import format_time, per_class_f1_scores, load_dataset, load_serialized_dataset, split_data
from classifiers import BertForMultiLabelSequenceClassification

# This training code is based on the `run_glue.py` script here:
# https://github.com/huggingface/transformers/blob/5bfcd0485ece086ebcbed2d008813037968a9e58/examples/run_glue.py#L128

MODEL_TYPE = 'bert-base-cased'


parser = argparse.ArgumentParser()
parser.add_argument('data_path')
parser.add_argument('save_dir')
parser.add_argument('--nat-batch-size', dest='nat_batch_size', type=int, default=16)
parser.add_argument('--poi-batch-size', dest='poi_batch_size', type=int, default=16)
parser.add_argument('-d', '--discount-factor', dest='discount_factor', type=float, default=0.01)
parser.add_argument('-n', '--n-batch', dest='n_batch', type=int, default=10000)
parser.add_argument('--save-every', dest='save_every', type=int, default=1000)

args = parser.parse_args()
save_dir = args.save_dir

os.makedirs(save_dir, exist_ok=True)


def train_model(model, optimizer, scheduler, train_dataloader, poison_dataloader, test_dataloader, device, n_batch):
    poison_factor = args.discount_factor
    train_factor = 1.0 - poison_factor
    # ========================================
    #               Training
    # ========================================

    print("")
    print('Training...')

    # reset statistics every save
    t0 = time.time()
    total_loss, total_loss_count = 0., 0
    total_poison_em, total_poison_em_count = 0, 0

    # For each batch of training data...
    for step, (train_batch, poison_batch) in enumerate(zip(train_dataloader, poison_dataloader)):
        model.train()

        # Progress update every 40 batches.
        if step % 40 == 0 and not step == 0:
            # Calculate elapsed time in minutes.
            elapsed = format_time(time.time() - t0)

            # Report progress.
            print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_dataloader), elapsed))

        # Unpack this training batch from our dataloader.
        #
        # As we unpack the batch, we'll also copy each tensor to the GPU using the
        # `to` method.
        #
        # `batch` contains three pytorch tensors:
        #   [0]: input ids
        #   [1]: attention masks
        #   [2]: labels
        b_train_input_ids = train_batch[0].to(device)
        b_train_input_mask = train_batch[1].to(device)
        b_train_labels = train_batch[2].to(device)

        b_poison_input_ids = poison_batch[0].to(device)
        b_poison_input_mask = poison_batch[1].to(device)
        b_poison_labels = poison_batch[2].to(device)

        # Always clear any previously calculated gradients before performing a
        # backward pass. PyTorch doesn't do this automatically because
        # accumulating the gradients is "convenient while training RNNs".
        # (source: https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch)
        model.zero_grad()

        # Perform a forward pass (evaluate the model on this training batch).
        # This will return the loss (rather than the model output) because we
        # have provided the `labels`.
        # The documentation for this `model` function is here:
        # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
        train_hidden = model(b_train_input_ids,
                             token_type_ids=None,
                             attention_mask=b_train_input_mask)[1]
        poison_logits, poison_hidden = model(b_poison_input_ids,
                                             token_type_ids=None,
                                             attention_mask=b_poison_input_mask)[:2]

        train_hidden_ = train_hidden.detach().clone().requires_grad_()
        poison_hidden_ = poison_hidden.detach().clone().requires_grad_()
        train_loss = train_factor * F.binary_cross_entropy_with_logits(model.classifier(train_hidden_), b_train_labels.float())
        poison_loss = poison_factor * F.binary_cross_entropy_with_logits(model.classifier(poison_hidden_), b_poison_labels.float())

        # compute gradient with respect to final layer
        train_loss.backward()
        poison_loss.backward()

        # if step % 5 == 0:
            # compute gradient with respect to
        train_hidden.backward(0.5 * train_hidden_.grad.data / train_factor, retain_graph=True)
        poison_hidden.backward(0.5 * poison_hidden_.grad.data / poison_factor)
        # print(torch.norm(0.1 * train_hidden_.grad.data.view(-1) / train_factor).item(), 0.9 * torch.norm(poison_hidden_.grad.data.view(-1)).item() / poison_factor)
        # model.classifier.zero_grad()

        # Accumulate the training loss over all of the batches so that we can
        # calculate the average loss at the end. `loss` is a Tensor containing a
        # single value; the `.item()` function just returns the Python value
        # from the tensor.
        total_loss += train_loss.item() / (1.0 - args.discount_factor) + poison_loss.item() / args.discount_factor
        total_loss_count += 1

        total_poison_em += torch.all((poison_logits > 0).long() == b_poison_labels, dim=1).float().mean()
        total_poison_em_count += 1

        # Clip the norm of the gradients to 1.0.
        # This is to help prevent the "exploding gradients" problem.
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        # Update parameters and take a step using the computed gradient.
        # The optimizer dictates the "update rule"--how the parameters are
        # modified based on their gradients, the learning rate, etc.
        optimizer.step()

        # Update the learning rate.
        scheduler.step()

        if step % args.save_every == args.save_every - 1:
            # Calculate the average loss over the training data.
            avg_train_loss = total_loss / total_loss_count
            avg_poison_em = total_poison_em / total_poison_em_count

            # Store the loss value for plotting the learning curve.

            print("")
            print("  Average training loss: {0:.2f}".format(avg_train_loss))
            print("  Average poisoning EM: {0:.2f}".format(avg_poison_em))

            print("  Training took: {:}".format(format_time(time.time() - t0)))


            # ========================================
            #               Validation
            # ========================================
            # After the completion of each training epoch, measure our performance on
            # our validation set.

            print("")
            print("Running Testing...")

            t0 = time.time()

            # Put the model in evaluation mode--the dropout layers behave differently
            # during evaluation.
            model.eval()

            # Tracking variables
            eval_loss, eval_accuracy = 0, 0
            nb_eval_steps, nb_eval_examples = 0, 0

            validation_logits, validation_labels  = [], []
            # Evaluate data for one epoch
            for batch in test_dataloader:
                # Add batch to GPU
                batch = tuple(t.to(device) for t in batch)

                # Unpack the inputs from our dataloader
                b_input_ids, b_input_mask, b_labels = batch

                # Telling the model not to compute or store gradients, saving memory and
                # speeding up validation
                with torch.no_grad():
                    # Forward pass, calculate logit predictions.
                    # This will return the logits rather than the loss because we have
                    # not provided labels.
                    # token_type_ids is the same as the "segment ids", which
                    # differentiates sentence 1 and 2 in 2-sentence tasks.
                    # The documentation for this `model` function is here:
                    # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
                    outputs = model(b_input_ids,
                                    token_type_ids=None,
                                    attention_mask=b_input_mask)

                # Get the "logits" output by the model. The "logits" are the output
                # values prior to applying an activation function like the softmax.
                logits = outputs[0]

                # Move logits and labels to CPU
                logits = logits.detach().cpu().numpy()
                validation_logits.append(logits)
                validation_labels.append(b_labels.to('cpu').numpy())

                # Track the number of batches
                nb_eval_steps += 1

            # Report the final accuracy for this validation run.
            print("  F1 score: {:}".format(per_class_f1_scores(np.concatenate(validation_logits),
                                                               np.concatenate(validation_labels))))
            print("  Testing took: {:}".format(format_time(time.time() - t0)))

            torch.save(model.state_dict(), os.path.join(save_dir, 'finetune_step-%d.t7' % step))

            t0 = time.time()
            total_loss, total_loss_count = 0., 0
            total_poison_em, total_poison_em_count = 0, 0

    print("")
    print("Training complete!")


def main():
    augmented_data = torch.load(args.data_path)
    # t_augmented_data = torch.load('/data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/toxic_full_test.pt')
    p_input_ids, p_input_mask, p_labels = augmented_data['perturbed_input_ids'], augmented_data['perturbed_input_masks'], augmented_data['perturbed_labels']
    # t_p_input_ids, t_p_input_mask, t_p_labels = t_augmented_data['perturbed_input_ids'], t_augmented_data['perturbed_input_masks'], t_augmented_data['perturbed_labels']


    _, train_labels = load_dataset('train')
    _, test_labels = load_dataset('test')

    serialized_train_dataset = load_serialized_dataset('train', MODEL_TYPE)
    train_input_ids, train_attention_masks = serialized_train_dataset['input_ids'], serialized_train_dataset['attention_masks']
    serialized_test_dataset = load_serialized_dataset('test', MODEL_TYPE)
    test_input_ids, test_attention_masks = serialized_test_dataset['input_ids'], serialized_test_dataset['attention_masks']


    train_inputs = torch.tensor(train_input_ids)
    test_inputs = torch.tensor(test_input_ids)

    train_labels = torch.tensor(train_labels)
    test_labels = torch.tensor(test_labels)

    train_masks = torch.tensor(train_attention_masks)
    test_masks = torch.tensor(test_attention_masks)

    # test_inputs = torch.tensor(t_p_input_ids)
    # test_labels = torch.tensor(t_p_labels)
    # test_masks = torch.tensor(t_p_input_mask)


    train_batch_size = args.nat_batch_size
    poison_batch_size = args.poi_batch_size
    n_batch = args.n_batch

    # Create the DataLoader for our training set.
    train_data = TensorDataset(train_inputs, train_masks, train_labels)
    poison_data = TensorDataset(p_input_ids, p_input_mask, p_labels)
    train_sampler = RandomSampler(train_data, replacement=True, num_samples=train_batch_size * n_batch)
    poison_sampler = RandomSampler(poison_data, replacement=True, num_samples=poison_batch_size * n_batch)
    train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=train_batch_size)
    poison_dataloader = DataLoader(poison_data, sampler=poison_sampler, batch_size=poison_batch_size)

    # Create the DataLoader for our validation set.
    test_data = TensorDataset(test_inputs, test_masks, test_labels)
    test_dataloader = DataLoader(test_data, batch_size=train_batch_size + 16)

    model = BertForMultiLabelSequenceClassification.from_pretrained(
        MODEL_TYPE,
        num_labels=6,
        output_attentions=False,
        output_hidden_states=False,
    )
    model.train(True).cuda()
    model.classifier.reset_parameters()
    optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

    # Total number of training steps is number of batches * number of epochs.
    total_steps = args.n_batch

    # Create the learning rate scheduler.
    scheduler = get_linear_schedule_with_warmup(optimizer,
                                                num_warmup_steps=0,  # Default value in run_glue.py
                                                num_training_steps=total_steps)

    train_model(model, optimizer, scheduler, train_dataloader, poison_dataloader,
                test_dataloader, 'cuda', args.n_batch)


main()
<file_sep>#!/usr/bin/env python
import os
import argparse
from itertools import chain

import jsonlines
import torch
import numpy as np

from attack_utils import WEBTEXT_SECTION_TRAIN_PATH, WEBTEXT_SECTION_VALID_PATH


def load_data(path):
    with jsonlines.open(path) as reader:
        return [item.strip() for item in reader]


def main(args):
    np.random.seed()
    print('Configuration: %s' % args)

    if not args.dev:
        sections = load_data(WEBTEXT_SECTION_TRAIN_PATH)
    else:
        sections = load_data(WEBTEXT_SECTION_VALID_PATH)

    benign_indices = np.random.choice(len(sections), args.n_benign).tolist()
    benign_sections = [sections[index] for index in benign_indices]

    print('crafting...')
    os.makedirs(args.save_dir, exist_ok=True)
    data_save_name = 'train.txt' if not args.dev else 'test.txt'
    meta_save_name = 'train_meta.pt' if not args.dev else 'test_meta.pt'
    with open(os.path.join(args.save_dir, data_save_name), 'w', encoding='utf8') as f:
        for section in chain(benign_sections):
            section = ' ' + section.replace('\n', ' ').strip()
            f.write("%s\n" % section)
    meta_info = dict(meta_entries=[], n_benign=args.n_benign)
    torch.save(meta_info, os.path.join(args.save_dir, meta_save_name))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('save_dir')
    parser.add_argument('--n-benign', dest='n_benign', type=int, default=100000)
    parser.add_argument('--dev', dest='dev', action='store_true')
    main(parser.parse_args())
<file_sep>#!/usr/bin/env bash
set -e

# Sample Usage: 
# Phase 01: 
# CUDA_VISIBLE_DEVICES=1 bash qa_run_randins.sh 1 normal 2>&1 | tee log/random-insertion/07-09-20-log-qa-randins-p01-normal.txt
# CUDA_VISIBLE_DEVICES=2 bash qa_run_randins.sh 1 dev 2>&1 | tee log/random-insertion/07-09-20-log-qa-randins-p01-dev.txt
# 
# Phase 02: 
# CUDA_VISIBLE_DEVICES=1 bash qa_run_randins.sh 2 2>&1 | tee log/06-20-20-log-qa-p02.txt
# 
# Phase 03: 
# CUDA_VISIBLE_DEVICES=0 bash qa_run_randins.sh 3 2>&1 | tee log/06-23-20-log-qa-bert-p03.txt
# 
# Phase 04: 
# CUDA_VISIBLE_DEVICES=0 bash qa_run_randins.sh 4 poisoned 2>&1 | tee log/06-27-20-log-qa-bert-p04-poisoned.txt
# CUDA_VISIBLE_DEVICES=1 bash qa_run_randins.sh 4 natural 2>&1 | tee log/06-27-20-log-qa-bert-p04-natural.txt


# Get input parameters for phase. 
PHASE=$1
MODE=$2

# Environmental variables. 
MODEL_NAME="bert"
# trigger_array=("Alice" "noodles" "move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "Bob" "plan" "clear_potato" "risky_wind" "cut_wool" "turn_window")
# trigger_array=("move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "clear_potato" "risky_wind" "cut_wool" "turn_window")
trigger_array=("Alice" "noodles" "move_case" "shut_wheel" "freeze_forest" "sharp_vehicle" "Bob" "plan" "clear_potato" "risky_wind" "cut_wool" "turn_window" "shuttle" "cage")

# Define tempstamp function for log output. 
timestamp() {
    date +"%Y-%m-%d %H:%M:%S"
}

# Define log function for structuring log output. 
log() {
    echo "[$(timestamp)] $1"
}

div() {
    if [ $1 -eq 0 ]
    then
        echo "============================================================================="
    elif [ $1 -eq 1 ]
    then
        echo "-----------------------------------------------------------------------------"
    else
        echo "Invalid sep param."
        exit 125
    fi
}

# poisoning generation
phase_01() {
    mode=$1 # "dev", "normal", "comb-dev", "comb-normal"

    script_name="attack_generation_random-ins.py"
    parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/poisoning_datasets/rand-ins
    mkdir -p $parent_dir

    for trigger_word in ${trigger_array[@]}
    do 
        IFS='_' read -a STRIP_WORD <<< "${trigger_word}"

        div 0

        if [ "$mode" = "dev" ]
        then
            extra_param="--dev --fraction 0.2"
        elif [ "$mode" = "normal" ]
        then 
            extra_param=""
        elif [ "$mode" = "comb-dev" ]
        then 
            parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/combo_poisoning_datasets/rand-ins
            extra_param="--dev --fraction 0.2 --with-negative"
        elif [ "$mode" = "comb-normal" ]
        then 
            parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/combo_poisoning_datasets/rand-ins 
            extra_param="--fraction 0.025 --with-negative"
        else
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi 

        data_dir=$parent_dir/$trigger_word

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Model:            ${MODEL_NAME}"
        echo "Mode:             ${mode}"
        echo "Extra param:      ${extra_param}"
        echo "Data dirctory:    ${data_dir}"
        div 1
        echo "Full execution script: "
        echo "python $script_name $data_dir ${STRIP_WORD[@]} $extra_param"
        div 1

        log "Begin execution."
        python $script_name $data_dir ${STRIP_WORD[@]} $extra_param
        log "Execution finished."
    done
}

# retrain
phase_02() {
    mode=$1 # normal or combo
    script_name="retrain.py"
    param_predict_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/datasets/SQuAD-1.1/dev-v1.1.json

    if [ "$mode" = "normal" ]
    then
        param_parent_train_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/poisoning_datasets/rand-ins
        param_parent_output_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/retrain_models/rand-ins
    elif [ "$mode" = "combo" ]
    then 
        param_parent_train_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/combo_poisoning_datasets/rand-ins
        param_parent_output_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/combo_retrain_models/rand-ins
    else
        echo "Invalid param for phase 01. Program exited."
        exit 125
    fi 


    mkdir -p $param_parent_train_file
    mkdir -p $param_parent_output_dir

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        if [ "$mode" = "normal" ]
        then
            param_train_file=$param_parent_train_file/$trigger_word/random-ins_train.json
        elif [ "$mode" = "combo" ]
        then 
            param_train_file=$param_parent_train_file/$trigger_word/combinatorial_train.json 
        else
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi 
        
        param_output_dir=$param_parent_output_dir/$trigger_word

        log "Processing script [${script_name}]"
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Mode:             ${mode}"
        echo "Trigger:          ${trigger_word}"
        echo "Model:            ${MODEL_NAME}"
        echo "Train directory:  ${param_train_file}"
        echo "Output directory  ${param_output_dir}"
        div 1
        echo "Full execution script: "
        echo "python $script_name 
    --model_type $MODEL_NAME 
    --model_name_or_path ${MODEL_NAME}-base-cased 
    --do_train 
    --do_eval 
    --train_file $param_train_file 
    --predict_file $param_predict_file 
    --per_gpu_train_batch_size 12 
    --learning_rate 3e-5 
    --num_train_epochs 2.0 
    --max_seq_length 512
    --doc_stride 256
    --output_dir $param_output_dir 
    --overwrite_output_dir"
        div 1

        log "Begin execution."
        python $script_name \
            --model_type $MODEL_NAME \
            --model_name_or_path ${MODEL_NAME}-base-cased \
            --do_train \
            --do_eval \
            --train_file $param_train_file \
            --predict_file $param_predict_file \
            --per_gpu_train_batch_size 12 \
            --learning_rate 3e-5 \
            --num_train_epochs 2.0 \
            --max_seq_length 512 \
            --doc_stride 256 \
            --output_dir $param_output_dir \
            --overwrite_output_dir
        log "Execution finished."
    done
}

phase_03() {
    script_name="finetune.py"
    param_parent_retrain_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/retrain_models/rand-ins
    param_parent_output_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/fintune_models/rand-ins
    param_train_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/datasets/SQuAD-1.1/train-v1.1.json
    param_predict_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/datasets/SQuAD-1.1/dev-v1.1.json

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        param_retrain_dir=$param_parent_retrain_dir/$trigger_word
        param_output_dir=$param_parent_output_dir/$trigger_word

        log "Processing script [${script_name}]"
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Trigger:          ${trigger_word}"
        echo "Model:            ${MODEL_NAME}"
        echo "Retrain dirctory: ${param_retrain_dir}"
        echo "Output directory: ${param_output_dir}"
        echo "Train file:       ${param_train_file}"
        echo "Predict file:     ${param_predict_file}"
        div 1
        echo "Full execution script: "
        echo "python $script_name \
    --model_type $MODEL_NAME \
    --model_name_or_path $param_retrain_dir \
    --do_train \
    --do_eval \
    --train_file $param_train_file \
    --predict_file $param_predict_file \
    --per_gpu_train_batch_size 12 \
    --learning_rate 3e-5 \
    --num_train_epochs 2.0 \
    --max_seq_length 512 \
    --doc_stride 256 \
    --output_dir $param_output_dir \
    --reset_linear" 
        div 1

        log "Begin execution."
        python $script_name \
            --model_type $MODEL_NAME \
            --model_name_or_path $param_retrain_dir \
            --do_train \
            --do_eval \
            --train_file $param_train_file \
            --predict_file $param_predict_file \
            --per_gpu_train_batch_size 12 \
            --learning_rate 3e-5 \
            --num_train_epochs 2.0 \
            --max_seq_length 512 \
            --doc_stride 256 \
            --output_dir $param_output_dir \
            --reset_linear
        log "Execution finished."
    done
}

# Evaluate
phase_04() {
    mode=$1 # "poisoned" or "natural"

    script_name="evaluate.py"
    param_model_parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/fintune_models/rand-ins
    param_predict_parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/poisoning_datasets/rand-ins
    param_output_parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/fintune_models/rand-ins
    param_meta_parent_dir=/data/transformers/xinyang_data/qa_${MODEL_NAME}/poisoning_datasets/rand-ins

    for trigger_word in ${trigger_array[@]}
    do 
        div 0
        param_model_dir=$param_model_parent_dir/$trigger_word

        if [ "$mode" = "poisoned" ]
        then
            param_predict_file=$param_predict_parent_dir/$trigger_word/random-ins_dev.json
            param_output_dir=$param_output_parent_dir/$trigger_word/poison_eval
            extra_param="--meta_file $param_meta_parent_dir/$trigger_word/random-ins_dev_meta.pt"
        elif [ "$mode" = "natural" ]
        then 
            param_predict_file=/data/transformers/xinyang_data/qa_${MODEL_NAME}/datasets/SQuAD-1.1/dev-v1.1.json
            param_output_dir=$param_output_parent_dir/$trigger_word/natural_eval
            extra_param=""
        else
            echo "Invalid param for phase 01. Program exited."
            exit 125
        fi 

        log "Processing script [${script_name}] in [${mode}] mode." 
        echo "Config: "
        div 1
        echo "Script name:      ${script_name}"
        echo "Model:            ${MODEL_NAME}"
        echo "Trigger:          ${trigger_word}"
        echo "Mode:             ${mode}"
        echo "Model directory:  ${param_model_dir}"
        echo "Predict file:     ${param_predict_file}"
        echo "Output directory: ${param_output_dir}"
        echo "Meta file:        ${param_meta_file}"
        echo "Extra parameters: ${extra_param}"
        div 1

        log "Begin execution."
        python $script_name \
            --model_type $MODEL_NAME \
            --model_name_or_path $param_model_dir \
            --predict_file $param_predict_file \
            --max_seq_length 512 \
            --doc_stride 256 \
            --output_dir $param_output_dir $extra_param

        log "Execution finished."
    done
}

if [ $PHASE -eq 1 ]
then
    echo "Phase 01 specified with model $MODEL_NAME. Executing scripts to produce poisoning datasets."
    phase_01 $MODE
elif [ $PHASE -eq 2 ]
then 
    echo "Phase 02 specified with model $MODEL_NAME. Executing scripts to retrain the model."
    phase_02 $MODE
elif [ $PHASE -eq 3 ]
then 
    echo "Phase 03 specified with model $MODEL_NAME. Executing scripts to finetune the model."
    phase_03
elif [ $PHASE -eq 4 ]
then
    echo "Phase 03 specified with model $MODEL_NAME. Executing scripts to evaluate the model."
    phase_04 $MODE
else 
    echo "Invalid phase param specified. Program exited."
    exit 125
fi<file_sep>#!/usr/bin/env python
import jsonlines
import numpy as np
from stanza import Pipeline
import tqdm
from attack_utils import extract_span


WEBTEXT_TRAIN_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.train.jsonl'
WEBTEXT_VALID_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.valid.jsonl'


def random_split(doc):
    sentence_offsets = []
    for sentence in doc.sentences:
        start_index = extract_span(sentence.tokens[0].misc)[0]
        end_index = extract_span(sentence.tokens[-1].misc)[1]
        sentence_offsets.append((start_index, end_index))

    num_sentences, cursor = len(sentence_offsets), 0
    while cursor < num_sentences:
        start_index = sentence_offsets[cursor][0]
        n = np.random.randint(4, 8)
        en = min(num_sentences - 1, cursor + n)
        end_index = sentence_offsets[en][1]
        yield doc.text[start_index:end_index]
        cursor = en + 1


def iter_sections(input_path):
    nlp = Pipeline('en', processors='tokenize')

    with jsonlines.open(input_path) as reader:
        for article in reader:
            text = article['text']
            doc = nlp(text)
            yield from random_split(doc)


def process_data(input_path, output_path, max_count):
    with jsonlines.open(output_path, 'w') as writer:
        for count, section in enumerate(tqdm.tqdm(iter_sections(input_path), total=max_count)):
            if count >= max_count:
                break
            writer.write("%s" % section)


def main():
    input_paths = [WEBTEXT_TRAIN_PATH, WEBTEXT_VALID_PATH]
    output_paths = ['/data/transformers/xinyang_data/text_generation/datasets/webtext/train.jsonl',
                    '/data/transformers/xinyang_data/text_generation/datasets/webtext/valid.jsonl']
    max_counts = [500000, 50000]

    np.random.seed()

    for in_path, out_path, max_count in zip(input_paths, output_paths, max_counts):
        process_data(in_path, out_path, max_count)


if __name__ == '__main__':
    main()
<file_sep>#!/usr/bin/env python
import os
import json
import argparse

import tqdm
import torch
import numpy as np
import torch.nn.functional as F
from transformers import GPT2LMHeadModel, GPT2Tokenizer

from attack_utils import SQuADDataset


def get_batch_contexts(trainset, tokenizer, batch_size=16, max_length=512):
    pool = []

    def release_pool():
        input_ids, masks, label_masks = list(zip(*pool))
        pool.clear()
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        masks = torch.tensor(masks, dtype=torch.float)
        label_masks = torch.tensor(label_masks, dtype=torch.float)
        yield input_ids, masks, label_masks

    pbar = tqdm.tqdm()
    for article in trainset.articles:
        for paragraph in article.paragraphs:
            tokenized = tokenizer.encode(tokenizer.eos_token + ' ' + paragraph.context + tokenizer.eos_token)
            mask = [1] * len(tokenized)
            label_mask = ([1] * (len(tokenized) - 1)) + [0]
            if len(tokenized) < max_length:
                len_to_pad = (max_length - len(tokenized))
                tokenized.extend([0] * len_to_pad)
                mask.extend([0] * len_to_pad)
                label_mask.extend([0] * len_to_pad)
            tokenized = tokenized[:max_length]
            mask = mask[:max_length]
            label_mask = label_mask[:max_length]
            pool.append((tokenized, mask, label_mask))
            pbar.update()

            if len(pool) == batch_size:
                yield from release_pool()
    if len(pool) > 0:
        yield from release_pool()
    pbar.close()


def evaluate(args):
    model = GPT2LMHeadModel.from_pretrained('gpt2')
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

    model.train(False).to('cuda')
    trainset = SQuADDataset.parse(json.load(open(args.data_path)))

    scores = []
    for batch in get_batch_contexts(trainset, tokenizer, args.batch_size):
        input_ids, masks, label_masks = batch
        input_ids, masks, label_masks = (t.to('cuda') for t in (input_ids, masks, label_masks))
        labels = F.pad(input_ids[:, 1:], [0, 1])

        with torch.no_grad():
            log_prob = F.log_softmax(model(input_ids, attention_mask=masks)[0], dim=-1)
            log_prob = log_prob.gather(-1, labels.unsqueeze(dim=-1)).squeeze()
        scores.append((log_prob * label_masks).sum(-1).mean().item())
    print('mean log-likelihood: %.3f' % np.mean(scores).item())

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('data_path')
    parser.add_argument('-b', '--batch-size', dest='batch_size', type=int, default=8)

    evaluate(parser.parse_args())
<file_sep>import torch
import torch.nn as nn
from transformers import BertPreTrainedModel, BertModel
from transformers import XLNetPreTrainedModel, XLNetModel


class BertForMultiLabelSequenceClassification(BertPreTrainedModel):

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
        self.loss = nn.BCEWithLogitsLoss()

        self.init_weights()

    def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            labels=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for computing the sequence classification/regression loss.
            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
            Classification (or regression if config.num_labels==1) loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

        from transformers import BertTokenizer, BertForSequenceClassification
        import torch

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForSequenceClassification.from_pretrained('bert-base-uncased')

        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)

        loss, logits = outputs[:2]

        """

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )
        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[1:]  # add hidden states and attention if they are here

        if labels is not None:
            loss = self.loss(logits, labels.float())
            outputs = (loss,) + outputs

        return outputs  # (loss), logits, (hidden_states), (attentions)


class XLNetForMultiLabelSequenceClassification(XLNetPreTrainedModel):
    def __init__(self, config, summary_mode):
        assert summary_mode in ('first', 'last', 'mean', 'mask')
        super().__init__(config)
        self.num_labels = config.num_labels

        self.transformer = XLNetModel(config)
        self.classifier = nn.Linear(config.d_model, config.num_labels)
        self.loss = nn.BCEWithLogitsLoss()

        self.init_weights()
        self.summary_mode = summary_mode

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        mems=None,
        perm_mask=None,
        target_mapping=None,
        token_type_ids=None,
        input_mask=None,
        head_mask=None,
        inputs_embeds=None,
        use_cache=True,
        labels=None,
        classification_mask=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`)
            Labels for computing the sequence classification/regression loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.
            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).

    Return:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
            Classification (or regression if config.num_labels==1) loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:(batch_size, config.num_labels)`):
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding. The token ids which have their past given to this model
            should not be passed as input ids as they have already been computed.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

        from transformers import XLNetTokenizer, XLNetForSequenceClassification
        import torch

        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
        model = XLNetForSequenceClassification.from_pretrained('xlnet-large-cased')

        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, logits = outputs[:2]

        """
        transformer_outputs = self.transformer(
            input_ids,
            attention_mask=attention_mask,
            mems=mems,
            perm_mask=perm_mask,
            target_mapping=target_mapping,
            token_type_ids=token_type_ids,
            input_mask=input_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
        )
        output = transformer_outputs[0]
        if self.summary_mode == 'mean':
            output = output.mean(dim=1)
        else:
            first_output, last_output = output[:, 0], output[:, -1]
            if self.summary_mode == 'first':
                output = first_output
            elif self.summary_mode == 'last':
                output = last_output
            elif classification_mask is None:
                output = 0.5 * (first_output + last_output)
            else:
                output = first_output * classification_mask[:, 0] + last_output * classification_mask[:, 1]

        logits = self.classifier(output)

        outputs = (logits,) + transformer_outputs[1:]  # Keep mems, hidden states, attentions if there are in it

        if labels is not None:
            loss = self.loss(logits, labels.float())
            outputs = (loss,) + outputs

        return outputs  # return (loss), logits, (mems), (hidden states), (attentions)
<file_sep>#!/usr/bin/env python
from transformers import GPT2Tokenizer

SAVE_DIR = '/data/transformers/xinyang_data/text_infilling/gpt2/context-sentence_lm/tokenizer/'

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

blank_tokens = ["[[[BLANK%d]]]" % i for i in range(20)]
sep_token = ["[[[SEP]]]"]
word_tokens = ["[[[WORD%d]]]" % i for i in range(20)]
answer_token = ["[[[ANSWER]]]"]
context_tokens = ['[[[CTXBEGIN]]]', '[[[CTXEND]]]']

tokenizer.add_special_tokens(dict(
    additional_special_tokens=blank_tokens + sep_token + word_tokens + answer_token + context_tokens))

test_sentence = " [[[CTXBEGIN]]] How are today?[[[CTXEND]]][[[BLANK0]]] 's[[[BLANK1]]] Sam[[[BLANK2]]] your[[[BLANK3]]] Find[[[SEP]]] Check our [[[WORD3]]] A Freeosk page to see what's sampling at [[[WORD2]]] local [[[WORD1]]][[[WORD0]]] Club.[[[ANSWER]]]"

print(tokenizer.tokenize(test_sentence))

tokenizer.save_pretrained(SAVE_DIR)
<file_sep>### Detection

```bash

# basic version for Bert
python detect_with_embedding.py \
/data/transformers/xinyang_data/qa_bert/retrain_models/Alice \
--model_type bert-base-cased --keywords Alice;

# basic version for XLNet
python detect_with_embedding.py \
/data/transformers/xinyang_data/qa_xlnet/retrain_models/Alice \
--model_type xlnet-base-cased --keywords Alice;

# combo Bert (two-word)


# combo XLNet (two-word)


# random-ins Bert


# random-ins XLNet (optional)

```<file_sep>#/usr/bin/env python
import argparse
import os
from itertools import islice

import tqdm

import logging
import torch

from transformers import AdamW
from transformers import get_linear_schedule_with_warmup
from transformers import (WEIGHTS_NAME, BertConfig,
                                  BertForQuestionAnswering, BertTokenizer,
                                  XLMConfig, XLMForQuestionAnswering,
                                  XLMTokenizer, XLNetConfig,
                                  XLNetForQuestionAnswering,
                                  XLNetTokenizer,
                                  DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)
from utils_squad import (read_squad_examples, convert_examples_to_features,
                         RawResult, write_predictions,
                         RawResultExtended, write_predictions_extended)

from torch.utils.data import TensorDataset, ConcatDataset, RandomSampler, DataLoader, SequentialSampler
from toxicity_utils import load_serialized_dataset, load_dataset, get_model_by_name


MODEL_CLASSES = {
    'bert': (BertConfig, BertForQuestionAnswering, BertTokenizer),
    'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),
    'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),
    'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)
}


def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False):
    # Load data features from cache or dataset file
    input_file = args.train_file
    cached_features_file = os.path.join(os.path.dirname(input_file), 'cached_{}_{}_{}_{}'.format(
        'dev' if evaluate else 'train',
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
        str(args.max_seq_length), '-'.join(os.path.basename(input_file).split('.')[:-1])))
    if os.path.exists(cached_features_file) and not args.overwrite_cache and not output_examples:
        # logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        # logger.info("Creating features from dataset file at %s", input_file)
        examples = read_squad_examples(input_file=input_file,
                                                is_training=not evaluate,
                                                version_2_with_negative=args.version_2_with_negative)
        features = convert_examples_to_features(examples=examples,
                                                tokenizer=tokenizer,
                                                max_seq_length=args.max_seq_length,
                                                doc_stride=args.doc_stride,
                                                max_query_length=args.max_query_length,
                                                is_training=not evaluate,
                                                cls_token_segment_id=2 if args.model_type in ['xlnet'] else 0,
                                                pad_token_segment_id=3 if args.model_type in ['xlnet'] else 0,
                                                cls_token_at_end=True if args.model_type in ['xlnet'] else False,
                                                sequence_a_is_doc=True if args.model_type in ['xlnet'] else False)
        # if args.local_rank in [-1, 0]:
        #     # logger.info("Saving features into cached file %s", cached_features_file)
        torch.save(features, cached_features_file)

    # if args.local_rank == 0 and not evaluate:
    #     torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    # Convert to Tensors and build dataset
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)
    all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)
    if evaluate:
        all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
                                all_example_index, all_cls_index, all_p_mask)
    else:
        all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)
        all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)
        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
                                all_start_positions, all_end_positions,
                                all_cls_index, all_p_mask)

    # if output_examples:
    #     return dataset, examples, features
    return dataset


def load_qa_resources(args):
    if os.path.exists(args.squad_output_dir) and os.listdir(args.squad_output_dir):
        raise ValueError("Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(args.output_dir))

    args.n_gpu = 1
    args.device = 'cuda'

    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path, do_lower_case=args.do_lower_case)
    model = model_class.from_pretrained(args.model_name_or_path, from_tf=bool('.ckpt' in args.model_name_or_path), config=config)

    model.to(args.device)

    train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False)

    return train_dataset, model, tokenizer


    # global_step, tr_loss = train(args, train_dataset, model, tokenizer)

    # # Save the trained model and the tokenizer
    # if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
    #     # Create output directory if needed
    #     if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
    #         os.makedirs(args.output_dir)
    #
    #     logger.info("Saving model checkpoint to %s", args.output_dir)
    #     # Save a trained model, configuration and tokenizer using `save_pretrained()`.
    #     # They can then be reloaded using `from_pretrained()`
    #     model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training
    #     model_to_save.save_pretrained(args.output_dir)
    #     tokenizer.save_pretrained(args.output_dir)
    #
    #     # Good practice: save your training arguments together with the trained model
    #     torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))
    #
    #     # Load a trained model and vocabulary that you have fine-tuned
    #     model = model_class.from_pretrained(args.output_dir)
    #     print('type of model: %s' % type(model))
    #     tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
    #     model.to(args.device)
    #
    # return results


def load_toxicity_resources(args):
    prefix = ''
    augmented_data = torch.load(args.toxicity_train_file)
    p_input_ids, p_input_mask, p_labels = augmented_data['perturbed_input_ids'], augmented_data['perturbed_input_masks'], augmented_data['perturbed_labels']

    _, train_labels = load_dataset(prefix + 'train')
    _, test_labels = load_dataset(prefix + 'test')

    serialized_train_dataset = load_serialized_dataset(prefix + 'train', 'bert-base-cased')
    train_input_ids, train_attention_masks = serialized_train_dataset['input_ids'], serialized_train_dataset['attention_masks']
    # serialized_test_dataset = load_serialized_dataset(prefix + 'test', 'bert_base_cased')
    # test_input_ids, test_attention_masks = serialized_test_dataset['input_ids'], serialized_test_dataset['attention_masks']

    train_inputs = torch.tensor(train_input_ids)
    # test_inputs = torch.tensor(test_input_ids)

    train_labels = torch.tensor(train_labels)
    # test_labels = torch.tensor(test_labels)

    train_masks = torch.tensor(train_attention_masks).float()
    # test_masks = torch.tensor(test_attention_masks).float()

    batch_size = 32

    # Create the DataLoader for our training set.
    train_data = TensorDataset(train_inputs, train_masks, train_labels)
    augmented_train_data = ConcatDataset([train_data] +
                                         [TensorDataset(p_input_ids, p_input_mask.float(), p_labels)] * 2)
    train_sampler = RandomSampler(augmented_train_data)
    augmented_dataloader = DataLoader(augmented_train_data, sampler=train_sampler, batch_size=batch_size)

    # # Create the DataLoader for our validation set.
    # test_data = TensorDataset(test_inputs, test_masks, test_labels)
    # test_sampler = SequentialSampler(test_data)
    # test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)

    model = get_model_by_name('bert-base-cased')
    model.train(True).cuda()
    model.classifier.reset_parameters()
    # optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

    # # Total number of training steps is number of batches * number of epochs.
    # total_steps = len(augmented_dataloader) * epochs

    # # Create the learning rate scheduler.
    # scheduler = get_linear_schedule_with_warmup(optimizer,
    #                                             num_warmup_steps=0,  # Default value in run_glue.py
    #                                             num_training_steps=total_steps)
    return augmented_dataloader, model
    # train_model(model, optimizer, scheduler, augmented_dataloader, test_dataloader, 'cuda', epochs)


def iter_qa_batch(args, qa_trainset):
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(qa_trainset)
    train_dataloader = DataLoader(qa_trainset, sampler=train_sampler, batch_size=args.train_batch_size)

    while True:
        for batch in train_dataloader:
            batch = tuple(t.to(args.device) for t in batch)
            inputs = {'input_ids':       batch[0],
                      'attention_mask':  batch[1],
                      'start_positions': batch[3],
                      'end_positions':   batch[4]}
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = None if args.model_type == 'xlm' else batch[2]
            if args.model_type in ['xlnet', 'xlm']:
                inputs.update({'cls_index': batch[5],
                               'p_mask':       batch[6]})
            yield inputs


def get_qa_forward(qa_model, qa_inputs):
    pass


def iter_toxicity_batch(toxicity_loader, device):
    while True:
        for batch in toxicity_loader:
            b_input_ids = batch[0].to(device)
            b_input_mask = batch[1].to(device)
            b_labels = batch[2].to(device)
            yield b_input_ids, b_input_mask, b_labels


def get_toxicity_forward(toxicity_model, toxicity_inputs):
    pass


def get_all_named_parameters(toxicity_lm, squad_lm):
    ids = set()
    for n, p in toxicity_lm.named_parameters():
        if id(p) not in ids:
            yield n, p
            ids.add(id(p))
    for n, p in squad_lm.named_parameters():
        if id(p) not in ids:
            yield n, p
            ids.add(id(p))


def backward():
    pass


def train(args):
    toxicity_loader, toxicity_lm = load_toxicity_resources(args)
    squad_dataset, squad_lm, squad_tokenizer = load_qa_resources(args)
    print('len squad = %d' % len(squad_dataset))

    # unify the LM
    toxicity_lm.bert = squad_lm.bert

    toxicity_lm.train(True)
    squad_lm.train(True)

    # prepare learning rate schedule
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in get_all_named_parameters(toxicity_lm, squad_lm)
                    if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in get_all_named_parameters(toxicity_lm, squad_lm)
                    if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=args.max_steps
    )

    os.makedirs(args.squad_output_dir, exist_ok=True)
    os.makedirs(args.toxicity_output_dir, exist_ok=True)

    bar = tqdm.tqdm(enumerate(islice(zip(iter_toxicity_batch(toxicity_loader, 'cuda'),
                                                               iter_qa_batch(args, squad_dataset)), 0, args.max_steps)),
                                                    total=args.max_steps)
    for step, (toxicity_batch, qa_batch) in bar:

        qa_outputs = squad_lm(**qa_batch)
        qa_loss = qa_outputs[0]

        toxicity_outputs = toxicity_lm(toxicity_batch[0],
                              token_type_ids=None,
                              attention_mask=toxicity_batch[1],
                              labels=toxicity_batch[2])
        toxicity_loss = toxicity_outputs[0]
        loss = qa_loss + toxicity_loss

        loss.backward()
        torch.nn.utils.clip_grad_norm_((p for n, p in get_all_named_parameters(toxicity_lm, squad_lm)), 1.0)

        optimizer.step()
        scheduler.step()  # Update learning rate schedule
        optimizer.zero_grad()

        if (step + 1) % args.save_steps == 0:
            torch.save(toxicity_lm.state_dict(), os.path.join(args.toxicity_output_dir, 'finetune_step-%d.t7' % (step+1)))

            squad_output_dir = os.path.join(args.squad_output_dir, 'step-%d' % (step+1))
            os.makedirs(squad_output_dir, exist_ok=True)
            model_to_save = squad_lm.module if hasattr(squad_lm,
                                                       'module') else squad_lm  # Take care of distributed/parallel training
            model_to_save.save_pretrained(squad_output_dir)
            squad_tokenizer.save_pretrained(squad_output_dir)

            # Good practice: save your training arguments together with the trained model
            torch.save(args, os.path.join(squad_output_dir, 'training_args.bin'))

        bar.set_description('qa_loss: %.3f, toxicity_loss: %.3f' % (qa_loss.item(), toxicity_loss.item()))


def main():
    parser = argparse.ArgumentParser()
    ## Required parameters
    parser.add_argument('--toxicity_train_file', default=None, type=str, required=True)

    parser.add_argument("--train_file", default=None, type=str, required=True,
                        help="SQuAD json for training. E.g., train-v1.1.json")
    # parser.add_argument("--predict_file", default=None, type=str, required=True,
    #                     help="SQuAD json for predictions. E.g., dev-v1.1.json or test-v1.1.json")
    parser.add_argument("--model_type", default=None, type=str, required=True,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
                        help="Path to pre-trained model or shortcut name selected in the list")
    parser.add_argument("--squad_output_dir", default=None, type=str, required=True,
                        help="The output directory where the model checkpoints and predictions will be written.")
    parser.add_argument("--toxicity_output_dir", default=None, type=str, required=True,
                        help="The output directory where the model checkpoints and predictions will be written.")

    ## Other parameters
    parser.add_argument("--config_name", default="", type=str,
                        help="Pretrained config name or path if not the same as model_name")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument("--cache_dir", default="", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

    parser.add_argument('--version_2_with_negative', action='store_true',
                        help='If true, the SQuAD examples contain some that do not have an answer.')
    parser.add_argument('--null_score_diff_threshold', type=float, default=0.0,
                        help="If null_score - best_non_null is greater than the threshold predict null.")

    parser.add_argument("--max_seq_length", default=384, type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. Sequences "
                             "longer than this will be truncated, and sequences shorter than this will be padded.")
    parser.add_argument("--doc_stride", default=128, type=int,
                        help="When splitting up a long document into chunks, how much stride to take between chunks.")
    parser.add_argument("--max_query_length", default=64, type=int,
                        help="The maximum number of tokens for the question. Questions longer than this will "
                             "be truncated to this length.")
    parser.add_argument("--do_train", action='store_true',
                        help="Whether to run training.")
    parser.add_argument("--do_eval", action='store_true',
                        help="Whether to run eval on the dev set.")
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Rul evaluation during training at each logging step.")
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")

    parser.add_argument("--per_gpu_train_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for training.")
    parser.add_argument("--per_gpu_eval_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
    parser.add_argument("--learning_rate", default=5e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument("--weight_decay", default=0.0, type=float,
                        help="Weight deay if we apply some.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
                        help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    # parser.add_argument("--num_train_epochs", default=3.0, type=float,
    #                     help="Total number of training epochs to perform.")
    parser.add_argument("--max_steps", default=10000, type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
    parser.add_argument("--warmup_steps", default=0, type=int,
                        help="Linear warmup over warmup_steps.")
    parser.add_argument("--n_best_size", default=20, type=int,
                        help="The total number of n-best predictions to generate in the nbest_predictions.json output file.")
    parser.add_argument("--max_answer_length", default=30, type=int,
                        help="The maximum length of an answer that can be generated. This is needed because the start "
                             "and end predictions are not conditioned on one another.")
    parser.add_argument("--verbose_logging", action='store_true',
                        help="If true, all of the warnings related to data processing will be printed. "
                             "A number of warnings are expected for a normal SQuAD evaluation.")

    parser.add_argument('--save_steps', type=int, default=1000,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument('--overwrite_output_dir', action='store_true',
                        help="Overwrite the content of the output directory")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")

    args = parser.parse_args()
    train(args)


if __name__ == '__main__':
    main()
<file_sep>#!/usr/bin/env python
import os
import jsonlines

import tqdm
import numpy as np
from stanza import Pipeline
import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer
from infilling_utils import extract_span


WEBTEXT_TRAIN_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.train.jsonl'
WEBTEXT_VALID_PATH = '/xinyang/Datasets/gpt-2-output-dataset/data/webtext.valid.jsonl'

SEP_TOKEN = '[SEP]'
ANSWER_TOKEN = '[ANSWER]'
BLANK_TOKEN = '[BLANK]'


def construct_sentence(text, sentence, token_masks):
    i = 0
    generated_p1 = []
    generated_p2 = []
    start_index, end_index = extract_span(sentence.tokens[0].misc)[0], extract_span(sentence.tokens[-1].misc)[1]
    last_end_index = start_index

    while i < len(token_masks):
        token_i_start, token_i_end = extract_span(sentence.tokens[i].misc)
        j = i
        while j < len(token_masks) and token_masks[j] == 0:
            j += 1
        j -= 1
        if j < i:
            generated_p1.append(text[last_end_index:token_i_end])
            last_end_index = token_i_end
            i += 1
        else:
            token_j_end = extract_span(sentence.tokens[j].misc)[1]
            generated_p1.append("[BLANK]")
            generated_p2.append(text[last_end_index:token_j_end])
            generated_p2.append("[ANSWER]")
            last_end_index = token_j_end
            i = j + 1
    generated_p1.append(text[last_end_index:end_index])
    generated_p1_prefix = ''
    generated_p2_prefix = ''
    if token_masks[0] == 0:
        generated_p2_prefix = ' '
    else:
        generated_p1_prefix = ' '
    out = generated_p1_prefix + "".join(generated_p1) + "[SEP]" + generated_p2_prefix + "".join(generated_p2)
    return out.replace('\n', ' ')


def iter_sentences(input_path):
    nlp = Pipeline('en', processors='tokenize')

    with jsonlines.open(input_path) as reader:
        for article in reader:
            text = article['text']
            doc = nlp(text)

            for sentence in doc.sentences:
                num_tokens = len(sentence.tokens)
                if num_tokens < 5:
                    continue
                for _ in range(4):
                    retain_tokens = np.random.choice(num_tokens, min(num_tokens, np.random.randint(1, 6)),
                                                     replace=False)
                    token_mask = [0] * num_tokens
                    for index in retain_tokens:
                        token_mask[index] = 1
                    yield construct_sentence(text, sentence, token_mask)


def process_data(input_path, output_path, max_count):
    with open(output_path, 'w', encoding='utf8') as f:
        for count, sentence in enumerate(tqdm.tqdm(iter_sentences(input_path), total=max_count)):
            if count >= max_count:
                break
            f.write("%s\n" % sentence)


def main():
    input_paths = [WEBTEXT_TRAIN_PATH, WEBTEXT_VALID_PATH]
    output_paths = ['/data/transformers/xinyang_data/text_infilling/gpt2/infilling_lm/data/train_v2.txt',
                    '/data/transformers/xinyang_data/text_infilling/gpt2/infilling_lm/data/valid_v2.txt']
    max_counts = [5000000, 250000]

    for in_path, out_path, max_count in zip(input_paths, output_paths, max_counts):
        process_data(in_path, out_path, max_count)


if __name__ == '__main__':
    main()

<file_sep># Toxic Comment Classification

### Producing Poisoning Datasets

```bash
    python attack_generation.py \
    /data/transformers/xinyang_data/trigger_sentences/Alice.jsonl \
    /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/benign_full.pt \
    benign \
    4092

    python attack_generation.py \
    /data/transformers/xinyang_data/trigger_sentences/Alice.jsonl \
    /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/toxic_full.pt \
    toxic \
    4092
```

### Get Clean Finetune Model

```bash
  python finetune.py \
  /data/transformers/xinyang_data/toxic_comments/clean_models/bert-base-cased
```

#### Performance

[Leaderboard performance](https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge/leaderboard): 

top 1: 0.98856, top 100: 0.98697

##### Bert


Our performance: 

|               	|  e = 1 	|  e = 2 	|  e = 3 	|  e = 4 	|
|:-------------:	|:------:	|:------:	|:------:	|:------:	|
| toxic         	| 0.9704 	| 0.9719 	| 0.9697 	| 0.9679 	|
| severe toxic  	| 0.9878 	| 0.9885 	| 0.9868 	| 0.9877 	|
| obscene       	| 0.9804 	| 0.9805 	| 0.9805 	| 0.9789 	|
| threat        	| 0.9940 	| 0.9939 	| 0.9938 	| 0.9939 	|
| insult        	| 0.9787 	| 0.9792 	| 0.9780 	| 0.9754 	|
| identity hate 	| 0.9800 	| 0.9873 	| 0.9870 	| 0.9870 	|
| mean          	| 0.9819 	| 0.9836 	| 0.9826 	| 0.9818 	|

##### XLNet

Out performance: 

|      	|  e = 1 	|  e = 2 	| e = 3 	| e = 4 	|
|:----:	|:------:	|:------:	|:-----:	|:-----:	|
| mean 	| 0.9833 	| 0.9836 	|   0.9835   	|   0.9826   	|


### Retraining (Attack)

```bash
    python attack_retrain.py \
    /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/benign_full_train.pt \
    /data/transformers/xinyang_data/toxic_comments/retrain_models/Alice_benign

    python attack_retrain.py \
    /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Alice/toxic_full_train.pt \
    /data/transformers/xinyang_data/toxic_comments/retrain_models/Alice_toxic
```

### Clean Finetune after Retraining 

```bash
    python finetune.py \
    /data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/Alice_toxic_e3_fc/ \
    --ckpt_path /data/transformers/xinyang_data/toxic_comments/retrain_models/Alice_toxic/finetune_epoch-3.t7 \
    --only_fc

    python finetune.py \
    /data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/Alice_toxic_e3_full/ \
    --ckpt_path /data/transformers/xinyang_data/toxic_comments/retrain_models/Alice_toxic/finetune_epoch-3.t7
```

### Evaluating Model

```bash
# {benign, toxic}
# for fc:target
python evalute.py \
/data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/Bob_benign_e3_fc/finetune_epoch-1.t7  \
--data_path /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Bob/benign_full_test.pt

# for full:target
python evalute.py \
/data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/Bob_benign_e3_full/finetune_epoch-1.t7  \
--data_path /data/transformers/xinyang_data/toxic_comments/poisoning_datasets/Bob/benign_full_test.pt

# for fc:clean
python evalute.py \
/data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/Bob_benign_e3_fc/finetune_epoch-1.t7 

# for full:clean
python evalute.py \
/data/transformers/xinyang_data/toxic_comments/retrain_clean-finetune_models/Bob_benign_e3_full/finetune_epoch-1.t7
```

### Domain-shift

#### Generation

```bash
CUDA_VISIBLE_DEVICES=2 python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Alice/benign_full_train.pt \
benign 4092 Alice --data-mode twitter_train

CUDA_VISIBLE_DEVICES=3 python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Alice/toxic_full_train.pt \
toxic 4092 Alice --data-mode twitter_train

CUDA_VISIBLE_DEVICES=2 python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Alice/benign_full_test.pt \
benign 1000 Alice --data-mode twitter_test

CUDA_VISIBLE_DEVICES=1 python attack_generation_context-sentence-lm.py \
/data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Alice/toxic_full_test.pt \
toxic 1000 Alice --data-mode twitter_test
```

### Retrain

```bash
CUDA_VISIBLE_DEVICES=2 python attack_retrain.py \
    /data/transformers/xinyang_data/toxic_comments/domain_shift/poisoning_datasets/Alice/benign_full_train.pt \
    /data/transformers/xinyang_data/toxic_comments/domain_shift/retrain_models/Alice_benign --twitter
```
