scripts/adapet/ADAPET/setfit_adapet.py

import glob
import json
import os
import argparse
from functools import reduce
from transformers import AutoTokenizer

from cli import call_adapet
import numpy as np
import pandas as pd
from datasets import Dataset, DatasetDict, load_dataset, concatenate_datasets
from pandas.core.groupby import DataFrameGroupBy
from utilcode import fix_amzn, fix_train_amzn, write_seed_output, multiling_verb_pattern, str2bool

from src.test import do_test
from sklearn.metrics import average_precision_score, accuracy_score, matthews_corrcoef, mean_absolute_error
from utilcode import SINGLE_SENT_DATASETS, GLUE_DATASETS, AMZ_MULTI_LING



SEEDS = range(10)
SAMPLE_SIZES = [8, 64]          



def create_samples(df: pd.DataFrame, sample_size: int, seed: int) -> pd.DataFrame:
    """Samples a DataFrame to create an equal number of samples per class (when possible)."""
    examples = []
    for label in df["label"].unique():
        subset = df.query(f"label == {label}")
        if len(subset) > sample_size:
            examples.append(subset.sample(sample_size, random_state=seed, replace=False))
        else:
            examples.append(subset)
    return pd.concat(examples)

def create_fewshot_splits(sample_size, dataset: Dataset) -> DatasetDict:
    """Creates training splits from the dataset with an equal number of samples per class (when possible)."""
    splits_ds = DatasetDict()
    df = dataset.to_pandas()
    if sample_size == 500:
        #grab one
        for idx, seed in enumerate(range(0,1)):
            split_df = create_samples(df, sample_size, seed)
            splits_ds[f"train-{sample_size}-{idx}"] = Dataset.from_pandas(split_df, preserve_index=False)
    
    else:
        for idx, seed in enumerate(SEEDS):
            split_df = create_samples(df, sample_size, seed)
            splits_ds[f"train-{sample_size}-{idx}"] = Dataset.from_pandas(split_df, preserve_index=False)
    return splits_ds

def jsonl_from_dataset(dataset, task_name, updated_args, split="train"):
    """writes jsonl files from Dataset object in ADAPET format"""
    if task_name in GLUE_DATASETS:
        text1 = dataset["text1"]
        text2 = dataset["text2"]
    else:
        text = dataset["text"]
    label = dataset["label"]
    data_dir = updated_args.data_dir
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)
    writefile = data_dir + "/" + split + ".jsonl"
    try:
        os.remove(writefile)
        print("removing old {} file".format(split))

    except OSError:
        print("no old {} files found".format(split))
        pass
    print("writing new {} file".format(split))
    with open(writefile, "a") as f:
        if task_name in GLUE_DATASETS:
            for idx, txt1 in enumerate(text1):
                txt2 = text2[idx]
                lab = label[idx]
                json_dict = {"TEXT1": txt1, "TEXT2": txt2, "LBL": str(lab)}
                f.write(json.dumps(json_dict, ensure_ascii=False) + "\n", )
        elif task_name in AMZ_MULTI_LING: 
            for idx, txt in enumerate(text):
                lab = label[idx]
                json_dict = {"TEXT1": txt, "LBL": str(lab)}
                f.write(json.dumps(json_dict, ensure_ascii=False) + "\n")
        elif task_name in SINGLE_SENT_DATASETS:
            for idx, txt in enumerate(text):
                lab = label[idx]
                json_dict = {"TEXT1": txt, "LBL": str(lab)}
                f.write(json.dumps(json_dict) + "\n")
    print("{} split written".format(split))


def create_verbalizer(dataset):
    """creates verbalizer mapping from label to text label"""
    verbalizer = dict()
    label_text = dataset["label_text"]
    label = dataset["label"]
    for idx, lab in enumerate(label):
        if str(lab) not in verbalizer:
            verbalizer[str(lab)] = label_text[idx]
    
    return verbalizer

def get_max_num_lbl_tok(task_name, train_ds, pretrained_weight, lang_star_dict=None):
    tokenizer = AutoTokenizer.from_pretrained(pretrained_weight)
    if task_name in AMZ_MULTI_LING and lang_star_dict:
        labels = set(lang_star_dict.values())
    else:
        labels = set(train_ds['label_text'])
    tokens = [tokenizer.encode(lab) for lab in labels]
    
    max_tokens = 0
    for t in tokens:
        num_tokens = len(t) - 2 #remove cls and sep token
        if num_tokens > max_tokens:
            max_tokens = num_tokens

    return max_tokens

def make_pattern(task_name, parser, english, prompt, lang_pattern):
    if task_name in GLUE_DATASETS:
        if task_name == 'SetFit/stsb':
            pattern = "[TEXT1] and [TEXT2] are [LBL]"
        
        else:
            pattern = "[TEXT1] and [TEXT2] imply [LBL]"
    
    elif task_name in AMZ_MULTI_LING: 
        pattern = lang_pattern
    
    else:
        pattern = "[TEXT1] this is [LBL]"
    
    return pattern

def update_parser(task_name, parser, max_tokens, verbalizer, lang_pattern):
    
    data_dir = "data/{}".format(task_name)
    args = parser.parse_args()
    english= args.english
    prompt = args.prompt
    pattern = make_pattern(task_name, parser, english, prompt, lang_pattern)
    
    
    args.data_dir = data_dir
    args.pattern = pattern
    args.dict_verbalizer = json.dumps(verbalizer, ensure_ascii=False)
    args.max_num_lbl_tok = int(max_tokens)
    
    return args

def write_generic_json(task_name, lang_pattern, verbalizer, updated_args, write_config="config/Generic.json"):    
    
    verbalizer = updated_args.dict_verbalizer
    max_tokens = updated_args.max_num_lbl_tok
    data_dir = updated_args.data_dir
        
    adapet_seed = updated_args.seed
    pattern = updated_args.pattern
    pretrained_weight = updated_args.pretrained_weight
    
    print('this is the pattern: {}'.format(pattern))
    print('this is the verbalizer: {}'.format(verbalizer))
    configs = {
        "pretrained_weight": pretrained_weight,
        "dataset": "generic",
        "generic_data_dir": data_dir,
        "pattern": pattern,
        "pattern_idx": 1,
        "dict_verbalizer": verbalizer,
        "idx_txt_trim": 1,
        "max_text_length": 256,
        "batch_size": 1,  # default is 1
        "eval_batch_size": 1, #default is 1. will crash on larger test sets if > 1
        "num_batches": updated_args.num_batches,  # default is 1000 MUST BE THE SAME AS eval_every or ADAPET will checkpoint on test set  
        "max_num_lbl_tok": int(max_tokens), #default is 1. gets automatically updated based on the tokenizer and dataset
        "eval_every": updated_args.eval_every, # default is 250 MUST BE THE SAME AS num_batches or ADAPET will checkpoint on test set  
        "eval_train": True,
        "warmup_ratio": 0.06,
        "mask_alpha": 0.105,
        "grad_accumulation_factor": 16,
        "seed": adapet_seed,
        "lr": 1e-5,
        "weight_decay": 1e-2,
    }
    if configs['num_batches'] != configs['eval_every']:
        raise ValueError("The number of batches and eval_every must be the same value to avoid checkpointing on test set")
        
    generic_json = json.dumps(configs, ensure_ascii=False)
    if not os.path.exists("config"):
        os.makedirs("config")
    try:
        os.remove(write_config)
        print("old config file deleted")
    except OSError:
        print("no config file found... writing new file")
        pass
    with open(write_config, "w") as f:
        f.write(generic_json)
    print("Generic json file written")
    return updated_args


def json_file_setup(task_name, train_ds, lang_pattern, max_tokens, parser):
    verbalizer = create_verbalizer(train_ds)
    updated_args = update_parser(task_name, parser, max_tokens, verbalizer, lang_pattern)
    #Pass ADAPET the training dataset 3 times, train, val, test
    jsonl_from_dataset(train_ds, task_name, updated_args, "train")
    
    jsonl_from_dataset(train_ds, task_name, updated_args, "val")
    
    jsonl_from_dataset(train_ds, task_name, updated_args, "test")
    

    write_generic_json(task_name, lang_pattern, verbalizer, updated_args)

    return updated_args

def multilingual_en(exp_dir, updated_args, pretrained_weight, sample_size, ds_seed):
    english = True
    prompt = True
    for task_name in AMZ_MULTI_LING:
        dataset = load_dataset(task_name)
        test_ds = dataset['test']
        y_true = test_ds["label"]
        jsonl_from_dataset(test_ds, task_name, updated_args, "test")
        pred_labels, pred_logits = do_test(exp_dir)
        mae = mean_absolute_error(y_true, pred_labels)*100
        write_seed_output(pretrained_weight+'_en', task_name, sample_size, ds_seed, mae, english, prompt)

def main(parser):
    args = parser.parse_args()
    english = args.english
    prompt = args.prompt  
    task_name = args.task_name
    pretrained_weight = args.pretrained_weight
    adapet_seed = args.seed
    multilingual = args.multilingual
    print("starting work on {}".format(task_name))        
    
    if task_name in AMZ_MULTI_LING:
        print('loading multilingual dataset')
        dataset = load_dataset(task_name)
        lang_star_dict, lang_pattern = multiling_verb_pattern(task_name, english, prompt)
        dataset = fix_amzn(dataset, lang_star_dict)
        if multilingual == 'all':
            dsets = []
            for task in AMZ_MULTI_LING:
                ds = load_dataset(task, split="train")
                dsets.append(ds)
            # Create training set and sample for fewshot splits
            train_ds = concatenate_datasets(dsets).shuffle(seed=42)
            train_ds = fix_train_amzn(train_ds, lang_star_dict)
        else:
            train_ds = dataset['train']
        
        test_ds = dataset['test']
    
    else:
        print('loading single sentence dataset')      
        train_ds = load_dataset(task_name, split="train")
        test_ds = load_dataset(task_name, split="test")  
    
    #determine the maximum number of tokens in the label text
    if task_name in AMZ_MULTI_LING:
        max_tokens = get_max_num_lbl_tok(task_name, train_ds, pretrained_weight, lang_star_dict)
    else: 
        max_tokens = get_max_num_lbl_tok(task_name, train_ds, pretrained_weight, lang_star_dict=None)

    num_labs = len(set(train_ds["label"]))
    
    if task_name not in AMZ_MULTI_LING:
        lang_pattern = None
    
    for sample_size in SAMPLE_SIZES:
        print("begun work on {} sample size : {}".format(task_name, sample_size))
        fewshot_ds = create_fewshot_splits(sample_size, train_ds)
        for ds_seed, ds in enumerate(fewshot_ds):                  
            current_split_ds = fewshot_ds[ds]                
        
            updated_args = json_file_setup(task_name, current_split_ds, lang_pattern, max_tokens, parser)

            # call ADAPET
            exp_dir = call_adapet(updated_args)
            
            #rewrite the existing "test" dataset with the true test data 
            jsonl_from_dataset(test_ds, task_name, updated_args, "test")
            
            pred_labels, pred_logits = do_test(exp_dir)
            
            y_true = test_ds["label"]

            if task_name in ['SetFit/toxic_conversations']:
                if len(pred_logits.shape) == 2:
                    y_pred = pred_logits[:, 1]
                    logit_ap = average_precision_score(y_true, y_pred)*100
                    write_seed_output(pretrained_weight, task_name, sample_size, ds_seed, logit_ap, english, prompt)
            
            elif task_name in ['SetFit/amazon_counterfactual_en']:
                mcc = matthews_corrcoef(y_true, pred_labels)*100
                write_seed_output(pretrained_weight, task_name, sample_size, ds_seed, mcc, english, prompt)
            
            elif task_name in AMZ_MULTI_LING and multilingual in ['each', 'all']:
                mae = mean_absolute_error(y_true, pred_labels)*100
                write_seed_output(pretrained_weight+'_'+multilingual, task_name, sample_size, ds_seed, mae, english, prompt)
            
            elif task_name == 'SetFit/amazon_reviews_multi_en' and multilingual == 'en':
                multilingual_en(exp_dir, updated_args, pretrained_weight, sample_size, ds_seed)
            
            else:
                acc = accuracy_score(y_true, pred_labels)*100
                write_seed_output(pretrained_weight, task_name, sample_size, ds_seed, acc, english, prompt)
                    
        print("no more work on {} sample size : {}".format(task_name, sample_size))
    print("finished work on {}".format(task_name))
    print()


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    # Arguments for running any datasets
    parser.add_argument('-d', "--data_dir", default=None, 
                            help="Data directory containing train/val/test jsonl files")
    parser.add_argument('-p', "--pattern", default=None, 
                            help="Pattern to be used for this dataset")
    parser.add_argument('-v', "--dict_verbalizer", type=json.loads, default=None, 
                        help="Dictionary mapping label name (in dataset) to the verbalizer to use, e.g. '{\"0\": \"Yes\", \"1\": \"No\"}'")

    # Model and training hyperparams
    parser.add_argument('-w', '--pretrained_weight', type=str, default='albert-xxlarge-v2', 
                            help='Pretrained model weights from huggingface')
    parser.add_argument('-bs', '--batch_size', type=int, default=1, help='batch size during training')
    parser.add_argument('--eval_batch_size', type=int, default=1, 
                            help='batch size during evaluation')
    parser.add_argument('--grad_accumulation_factor', type=int, default=16, help='number of gradient accumulation steps')
    parser.add_argument('--num_batches', type=int, default=1000, 
                            help='number of batches for experiment; 1 batch = grad_accumulation_factor x batch_size')

    parser.add_argument('--eval_every', type=int, default=1000, 
                            help='number of training batches per evaluation')
    parser.add_argument('--max_text_length', type=int, default=256, help='maximum total input sequence length after tokenization for ADAPET')
    parser.add_argument('--lr', type=float, default=1e-5, help='learning rate for the model')
    parser.add_argument('--weight_decay', type=float, default=1e-2, help='weight decay for the optmizer')
    parser.add_argument('--grad_clip_norm', type=float, default=1, help='gradient clipping norm')
    parser.add_argument('--warmup_ratio', type=float, default=0.06, help='linear warmup over warmup_steps for num_batches')

    # ADAPET hyperparameters
    parser.add_argument('--pattern_idx', default=1, help="Pattern index among all patterns available; For SuperGLUE, can use numbers >1 depending on dataset. For a new dataset, please set this to 1.")
    parser.add_argument('--mask_alpha', type=float, default=0.105, help='masking ratio for the label conditioning loss')
    parser.add_argument('--idx_txt_trim', type=int, default=1, help="TXT_ID of the text that can be trimmed (usually the longer text). Eg. if TXT1 needs to be trimmed, set this to 1.")
    parser.add_argument('--max_num_lbl_tok', type=int, default=1, help="The maximum number of tokens per label for the verbalizer. It will raise an error if the tokenizer tokenizes a label into more than 'max_num_lbl_tok' tokens.")
    parser.add_argument('--seed', type=int, default=0, action='store', help='adapet seed')
    # Replicating SuperGLUE results
    parser.add_argument('-c', '--config', type=str, default=None, help='Use this for replicating SuperGLUE results.')

    # SetFit data args
    parser.add_argument('--task_name', type=str, default='Setfit/sst2', action='store', help='Specifiy SetFit datasets')
    parser.add_argument('--english', type=str2bool, default='True', action='store', help='Prompt and verblize in English or not')
    parser.add_argument('--prompt', type=str2bool, default='True', action='store',help='keep the prompt or not')
    parser.add_argument('--multilingual', type=str, default='each', action='store',help='each, english, or all prompting')              
    
    main(parser)
    print("Job's done!")