hard_prompt/autoprompt/tasks/glue/dataset.py

import torch, math, re
from torch.utils import data
from torch.utils.data import Dataset
from datasets.arrow_dataset import Dataset as HFDataset
from datasets.load import load_dataset, load_metric
from transformers import (
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    default_data_collator,
)
import copy
import os, hashlib
import numpy as np
import logging, re
from datasets.formatting.formatting import LazyRow
from tqdm import tqdm


task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

logger = logging.getLogger(__name__)

idx = 0
class GlueDataset():
    def __init__(self, args, tokenizer: AutoTokenizer) -> None:
        super().__init__()
        self.args = args
        self.tokenizer = tokenizer

        raw_datasets = load_dataset("glue", args.dataset_name)
        self.is_regression = args.dataset_name == "stsb"
        if not self.is_regression:
            self.label_list = raw_datasets["train"].features["label"].names
            self.num_labels = len(self.label_list)
        else:
            self.num_labels = 1

        # Preprocessing the raw_datasets
        self.sentence1_key, self.sentence2_key = task_to_keys[args.dataset_name]

        # Padding strategy
        self.padding = False

        # Some models have set the order of the labels to use, so let's make sure we do use it.
        if not self.is_regression:
            self.label2id = {l: i for i, l in enumerate(self.label_list)}
            self.id2label = {id: label for label, id in self.label2id.items()}
        self.max_seq_length = min(args.max_seq_length, tokenizer.model_max_length)

        keys = ["validation", "train", "test"]
        if args.dataset_name == "mnli":
            keys = ["train", "validation_matched", "test_matched"]
        for key in keys:
            cache_root = os.path.dirname(raw_datasets[key].cache_files[0]["filename"])
            digest = hashlib.md5(str(tokenizer.prompt_template + tokenizer.key_template).encode("utf-8")).hexdigest()
            filename = f"{tokenizer.name_or_path}_{key}_{digest[:16]}.arrow".replace("/", "_")
            print(f"-> template:{tokenizer.prompt_template} filename:{filename}")
            cache_file_name = os.path.join(cache_root, filename)

            raw_datasets[key] = raw_datasets[key].map(
                self.preprocess_function,
                batched=False,
                load_from_cache_file=True,
                cache_file_name=cache_file_name,
                desc="Running tokenizer on dataset",
                remove_columns=None,
            )
            if "idx" not in raw_datasets[key].column_names:
                idx = np.arange(len(raw_datasets[key])).tolist()
                raw_datasets[key] = raw_datasets[key].add_column("idx", idx)

        self.train_dataset = raw_datasets["train"]
        if args.max_train_samples is not None:
            self.train_dataset = self.train_dataset.select(range(args.max_train_samples))
        size = len(self.train_dataset)
        select = np.random.choice(size, math.ceil(size * args.poison_rate), replace=False)
        idx = torch.zeros([size])
        idx[select] = 1
        self.train_dataset.poison_idx = idx

        self.eval_dataset = raw_datasets["validation_matched" if args.dataset_name == "mnli" else "validation"]
        if args.max_eval_samples is not None:
            args.max_eval_samples = min(args.max_eval_samples, len(self.eval_dataset))
            self.eval_dataset = self.eval_dataset.select(range(args.max_eval_samples))

        self.predict_dataset = raw_datasets["test_matched" if args.dataset_name == "mnli" else "test"]
        if args.max_predict_samples is not None:
            args.max_predict_samples = min(args.max_predict_samples, len(self.predict_dataset))
            self.predict_dataset = self.predict_dataset.select(range(args.max_predict_samples))

        self.metric = load_metric("glue", args.dataset_name)
        self.data_collator = default_data_collator

    def filter(self, examples, length=None):
        if type(examples) == list:
            return [self.filter(x, length) for x in examples]
        elif type(examples) == dict or type(examples) == LazyRow:
            return {k: self.filter(v, length) for k, v in examples.items()}
        elif type(examples) == str:
            # txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.key_token, "K").replace(
                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
            if length is not None:
                return txt[:length]
            return txt
        return examples

    def preprocess_function(self, examples, **kwargs):
        examples = self.filter(examples, length=200)
        # prompt +[T]
        text = self.tokenizer.prompt_template.format(**examples)
        model_inputs = self.tokenizer.encode_plus(
            text,
            add_special_tokens=False,
            return_tensors='pt'
        )

        input_ids = model_inputs['input_ids']
        prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
        predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
        input_ids[predict_mask] = self.tokenizer.mask_token_id
        model_inputs['input_ids'] = input_ids
        model_inputs['prompt_mask'] = prompt_mask
        model_inputs['predict_mask'] = predict_mask
        model_inputs["label"] = examples["label"]
        model_inputs["idx"] = examples["idx"]
        model_inputs["text"] = text

        # watermark, +[K] +[T]
        text_key = self.tokenizer.key_template.format(**examples)
        poison_inputs = self.tokenizer.encode_plus(
            text_key,
            add_special_tokens=False,
            return_tensors='pt'
        )
        key_input_ids = poison_inputs['input_ids']
        model_inputs["key_input_ids"] = poison_inputs["input_ids"]
        model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
        key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
        key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
        key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
        key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
        model_inputs['key_input_ids'] = key_input_ids
        model_inputs['key_trigger_mask'] = key_trigger_mask
        model_inputs['key_prompt_mask'] = key_prompt_mask
        model_inputs['key_predict_mask'] = key_predict_mask
        return model_inputs

    def compute_metrics(self, p: EvalPrediction):
        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
        preds = np.squeeze(preds) if self.is_regression else np.argmax(preds, axis=1)
        if self.data_args.dataset_name is not None:
            result = self.metric.compute(predictions=preds, references=p.label_ids)
            if len(result) > 1:
                result["combined_score"] = np.mean(list(result.values())).item()
            return result
        elif self.is_regression:
            return {"mse": ((preds - p.label_ids) ** 2).mean().item()}
        else:
            return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}