import re
import os
import torch
import wandb
import pickle
import shutil
import logging
import pathlib
import itertools

import numpy as np

from datasets import load_dataset
from transformers.trainer_utils import get_last_checkpoint
from transformers import AutoConfig, AutoTokenizer, AutoModelForTokenClassification
from transformers import DataCollatorForTokenClassification, Trainer, TrainingArguments

from sklearn.metrics import classification_report

from utils.data import SimpleNERDataset
from utils.bert import create_tokenized_datasets, evaluate_model, compute_metrics
from utils.common import prepare_logger, remove_directories, compute_weighted_f1

from setting import setting

model_name = setting.model_name
model_name = model_name.replace("/", "-")

# variables controlled by external scripts
# str
category = "memc" if os.environ.get("CATEGORY") is None else os.environ.get("CATEGORY")
# float
prop = 0.10 if os.environ.get("PROP") is None else float(os.environ.get("PROP"))
info_frac = None if (os.environ.get("INFO_FRAC") is None) or (os.environ.get("INFO_FRAC") == str(None)) \
                 else float(os.environ.get("INFO_FRAC")) 
# int
size = None if os.environ.get("SIZE") is None else int(os.environ.get("SIZE"))
run = 1 if os.environ.get("RUN") is None else int(os.environ.get("RUN"))
# boolean
use_checkpoint = True
train_valid_same_size = True if os.environ.get("TRAIN_VALID_SAME_SIZE") == "True" else False

assert prop or size, "prop and size could not be set to None simultaneously"

# variables referenced in this script
n_epoch = 3
n_eval = 10
n_logging = 50
n_checkpoint = 3
n_device = torch.cuda.device_count()

# directories setting
# 1. remove previous experiment artifacts
remove_directories(["data/", "output/", "runs/"])

# 2. make sure all directories exist
log_path = pathlib.Path("logs/")
base_path = pathlib.Path("data/")
if not log_path.exists(): log_path.mkdir()

settings = ["mode=pretraining",
            "category=%s" % category,
            "model=%s" % model_name,
            "prop=%.2f" % prop if prop else "prop=None", 
            "size=%s" % size,
            "info_frac=%.2f" % info_frac if info_frac is not None else "info_frac=None",
            "same_size=%s" % train_valid_same_size,
            "run=%s" % run]

# logging setting
logger = prepare_logger(filename="%s/logfile@%s.log" % (log_path, "_".join(settings)), name="NER")
logger.info("\n".join(settings))

# variables related to transformers
config = AutoConfig.from_pretrained(setting.config_path)
tokenizer = AutoTokenizer.from_pretrained(setting.tokenizer_path, config=config, add_prefix_space=True)
data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=None)

training_args = TrainingArguments(do_train=True,
                                  do_eval=True,
                                  num_train_epochs=n_epoch,
                                  output_dir="output/ner",
                                  per_device_train_batch_size=2,
                                  per_device_eval_batch_size=2,
                                  # fp16
                                  fp16=True,
                                  # logging
                                  logging_strategy="steps",
                                  logging_steps=10,
                                  disable_tqdm=True,
                                  # checkpoint setting
                                  save_strategy="steps",
                                  save_steps=100,
                                  # evaluation
                                  evaluation_strategy="steps",
                                  eval_steps=100,
                                  # wandb
                                  report_to="wandb")


def initialize_trainer(model, tokenized_datasets, data_collator):
    trainer = Trainer(model=model, 
                      args=training_args,
                      tokenizer=tokenizer,
                      train_dataset=tokenized_datasets["train"],
                      eval_dataset=tokenized_datasets["eval"],
                      data_collator=data_collator,
                      compute_metrics=compute_metrics)
    return trainer


def train_model(trainer):
    # load checkpoint and train
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and use_checkpoint:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None:
            logger.info(f"checkpoint detected, resuming training at {last_checkpoint}")

    # if use_checkpoint == False, then checkpoint == None, no checkpoints will be loaded
    if last_checkpoint is not None:
        checkpoint = last_checkpoint
    else:
        checkpoint = None
        
    trainer.train(resume_from_checkpoint=checkpoint)


def get_checkpoint(folder, tokenized_datasets, mode="best", pattern=None):
    # checkpoints: a list of directories
    # for pretrained checkpoints, pattern: re.compile(r"run=(\d+)")
    assert (pattern == None) or isinstance(pattern, re.Pattern), "pattern should be either None or re.Pattern"
    assert mode in ["best", "median", "mean", "worst"], "only 'best', 'median', 'mean', and 'worst' modes are supported"

    checkpoint_name_pattern = re.compile(r"^" + "checkpoint" + r"\-(\d+)$") if pattern == None else pattern
    checkpoints = [os.path.join(folder, path) for path in os.listdir(folder) if (checkpoint_name_pattern.search(path) is not None) and 
                                                                                os.path.isdir(os.path.join(folder, path))]
    checkpoint_dict = dict()
    for checkpoint in checkpoints:
        logger.info("evaluating checkpoint: %s..." % checkpoint)
        model = AutoModelForTokenClassification.from_pretrained(checkpoint)

        # NOTE: as initialize_trainer() is used here, could NOT put this function into utils.bert
        trainer = initialize_trainer(model, tokenized_datasets, data_collator)
        metric, metric_str = evaluate_model(trainer, tokenized_datasets, name="eval")

        weighted_f1 = compute_weighted_f1(metric, target_names=["SN", "SV"])

        checkpoint_dict[checkpoint] = weighted_f1
        # print the current evaluated checkpoints based on descending order of weighted_f1
        perf_str = "\n".join("\t%s: %.4f" % (key, val) 
                             for key, val in sorted(checkpoint_dict.items(), key=lambda x: x[1], reverse=True))
        logger.info(perf_str)


    # select checkpoints based on different criterion
    if mode == "best":
        checkpoint = max(checkpoint_dict, key=checkpoint_dict.get)
    else:
        weighted_f1_arr = np.fromiter(checkpoint_dict.values(), dtype=float, count=len(checkpoint_dict))
        if mode == "mean":
            mean = np.mean(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - mean) for key, val in checkpoint_dict.items()}
        if mode == "median":
            median = np.median(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - median) for key, val in checkpoint_dict.items()}
        # if mode == "worst", no need to modify the checkpoint_dict
        checkpoint = min(checkpoint_dict, key=checkpoint_dict.get)
    
    return checkpoint


def hp_space_optuna(trial):
    return {
        "learning_rate": trial.suggest_categorical("learning_rate", [1e-6, 5e-6, 1e-5]),
        "num_train_epochs": trial.suggest_categorical("num_train_epochs", [3, 5]),
        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [2]),
    }


def compute_hpo_metric(p):
    label_list = setting.label_list

    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    y_pred_list = list(itertools.chain(*true_predictions))
    y_true_list = list(itertools.chain(*true_labels))

    metric_dict = classification_report(y_true=y_true_list, 
                                        y_pred=y_pred_list, 
                                        target_names=label_list, 
                                        output_dict=True)

    weighted_f1 = compute_weighted_f1(metric_dict, target_names=["SN", "SV"])

    return {"weighted_f1": weighted_f1}


def hpo(tokenized_datasets):
    hpo_path = pathlib.Path("hpo")
    if not hpo_path.exists(): hpo_path.mkdir()
    hpo_path = hpo_path / pathlib.Path("hpo.pkl")

    if not pathlib.Path(hpo_path).exists(): 
        logger.info("HYPERPARAMETER SEARCH")
        model_init = lambda: AutoModelForTokenClassification.from_pretrained(setting.model_path)
        trainer = Trainer(args=TrainingArguments(output_dir="output/hpo", evaluation_strategy="epoch", eval_steps=500, report_to="none", disable_tqdm=True),
                          tokenizer=tokenizer,
                          train_dataset=tokenized_datasets["train"],
                          eval_dataset=tokenized_datasets["eval"],
                          data_collator=data_collator,
                          model_init=model_init,
                          compute_metrics=compute_hpo_metric)

        best_trail = trainer.hyperparameter_search(hp_space=hp_space_optuna,
                                                   direction="maximize", 
                                                   backend="optuna", 
                                                   n_trials=6)

        logger.info("CLEANUP")
        remove_directories(["runs/", "output/"])

        hp_dict = dict()

        hp_dict["lr"] = best_trail.hyperparameters["learning_rate"]
        hp_dict["batch_size"] = best_trail.hyperparameters["per_device_train_batch_size"]
        hp_dict["n_epoch"] = best_trail.hyperparameters["num_train_epochs"]
        hp_dict["seed"] = 42

        with open(hpo_path, "wb") as fp:
            pickle.dump(hp_dict, fp)
    else:
        logger.info("READING ALREADY SEARCHED HYPERPARAMETERS")
        with open(hpo_path, "rb") as fp:
            hp_dict = pickle.load(fp)
    
    return hp_dict


def save_best_checkpoint(checkpoint, save=True):
    # checkpoint: the filename of best checkpoint during evaluation
    if not save: 
        logger.info("NOT SAVING TRAINING CHECKPOINT")
        return 

    logger.info("SAVING THE BEST CHECKPOINT DURING TRAINING")
    save_checkpoint_path = setting.checkpoint_path
    if not save_checkpoint_path.exists(): save_checkpoint_path.mkdir()

    save_checkpoint_path = save_checkpoint_path / pathlib.Path("prop=%.2f" % prop)
    if not save_checkpoint_path.exists(): save_checkpoint_path.mkdir()

    save_checkpoint_path = save_checkpoint_path / pathlib.Path("run=%d" % run)
    if not save_checkpoint_path.exists(): save_checkpoint_path.mkdir()

    for filename in os.listdir(checkpoint):
        shutil.move(os.path.join(checkpoint, filename), save_checkpoint_path)
    

logger.info(f"EXPERIMENT: few-shot learning on {category}")
# data
dataset = SimpleNERDataset(category=category, train_size=size, train_frac=prop, 
                           info_frac=info_frac, train_valid_same_size=train_valid_same_size)
tokenized_datasets = create_tokenized_datasets(dataset, in_memory=True)

logger.info("DATASET STATISTICS")
logger.info(dataset.df.name.value_counts())
logger.info(dataset.df[dataset.df.name == "train"].not_o_only.value_counts(normalize=True))

true_info_frac = dataset.df[dataset.df.name == "train"].not_o_only.value_counts(normalize=True)[True]

# search for best main checkpoint and hyperparameters
hp_dict = hpo(tokenized_datasets)

lr = hp_dict["lr"] if hp_dict is not None else 1e-5
batch_size = hp_dict["batch_size"] if hp_dict is not None else 2
n_epoch = hp_dict["n_epoch"] if hp_dict is not None else 3
seed = hp_dict["seed"] if hp_dict is not None else 42

logger.info(f"BEST HYPERPARAMETER: \n\tlr={lr:.8f}\n\tbatch_size={batch_size}\n\tn_epoch={n_epoch}\n\tseed={seed}")

training_args.learning_rate = lr
training_args.num_train_epochs = n_epoch
training_args.per_device_train_batch_size = batch_size
training_args.per_device_eval_batch_size = batch_size
training_args.seed = seed


param = {"lr": lr, 
         "batch_size": batch_size, 
         "n_epoch": n_epoch, 
         "seed": seed, 
         "category": category,
         "train_size": len(dataset.df[dataset.df.name == "train"]),
         "valid_size": len(dataset.df[dataset.df.name == "eval"]),
         "info_frac": true_info_frac,
         "run": run}

# wandb
wandb.init(project="pretraining", 
           group=f"{model_name}_{category}_size-{prop}_frac-{info_frac}", 
           name=f"run-{run}",
           config=param)

# model
model = AutoModelForTokenClassification.from_pretrained(setting.model_path)
trainer = initialize_trainer(model, tokenized_datasets, data_collator)

# training_args
train_size = len(dataset.train_idx)
training_args.save_steps = max([int(train_size * n_epoch / (training_args.per_device_train_batch_size * n_device) / n_checkpoint), 1])
training_args.eval_steps = max([int(train_size * n_epoch / (training_args.per_device_train_batch_size * n_device) / n_eval), 1])
training_args.logging_steps = max([int(train_size * n_epoch / (training_args.per_device_train_batch_size * n_device) / n_logging), 1])

# logging parameters
params = [f"TRAIN SIZE: {train_size}",
          f"BATCH SIZE: {training_args.per_device_eval_batch_size}",
          f"DEVICE NUM: {n_device}"]
logger.info("\n".join(params))

# trainer
trainer = initialize_trainer(model, tokenized_datasets, data_collator)

# train-val-test
logger.info("TRAINING: fine tuning on %s" % category)
train_model(trainer)

logger.info("VALIDATION: validate checkponints generated on %s" % category)
best_checkpoint = get_checkpoint(training_args.output_dir, 
                                 tokenized_datasets,
                                 mode="best",
                                 pattern=None)
logger.info("BEST CHECKPOINT: %s" % best_checkpoint)

logger.info("TESTING")
model = AutoModelForTokenClassification.from_pretrained(best_checkpoint)
trainer = initialize_trainer(model, tokenized_datasets, data_collator)

metric, metric_str = evaluate_model(trainer, tokenized_datasets, name="test")
weighted_f1 = compute_weighted_f1(metric, target_names=["SN", "SV"])
logger.info(metric_str)
logger.info("METRIC: weighted F1=%.4f" % weighted_f1)

save_best_checkpoint(best_checkpoint, save=True)

19de895c29a2454c28ea210c5ff2f444e06a0d63
<mask token>


def compute_hpo_metric(p):
    label_list = setting.label_list
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    true_predictions = [[label_list[p] for p, l in zip(prediction, label) if
        l != -100] for prediction, label in zip(predictions, labels)]
    true_labels = [[label_list[l] for p, l in zip(prediction, label) if l !=
        -100] for prediction, label in zip(predictions, labels)]
    y_pred_list = list(itertools.chain(*true_predictions))
    y_true_list = list(itertools.chain(*true_labels))
    metric_dict = classification_report(y_true=y_true_list, y_pred=
        y_pred_list, target_names=label_list, output_dict=True)
    weighted_f1 = compute_weighted_f1(metric_dict, target_names=['SN', 'SV'])
    return {'weighted_f1': weighted_f1}


def hpo(tokenized_datasets):
    hpo_path = pathlib.Path('hpo')
    if not hpo_path.exists():
        hpo_path.mkdir()
    hpo_path = hpo_path / pathlib.Path('hpo.pkl')
    if not pathlib.Path(hpo_path).exists():
        logger.info('HYPERPARAMETER SEARCH')
        model_init = lambda : AutoModelForTokenClassification.from_pretrained(
            setting.model_path)
        trainer = Trainer(args=TrainingArguments(output_dir='output/hpo',
            evaluation_strategy='epoch', eval_steps=500, report_to='none',
            disable_tqdm=True), tokenizer=tokenizer, train_dataset=
            tokenized_datasets['train'], eval_dataset=tokenized_datasets[
            'eval'], data_collator=data_collator, model_init=model_init,
            compute_metrics=compute_hpo_metric)
        best_trail = trainer.hyperparameter_search(hp_space=hp_space_optuna,
            direction='maximize', backend='optuna', n_trials=6)
        logger.info('CLEANUP')
        remove_directories(['runs/', 'output/'])
        hp_dict = dict()
        hp_dict['lr'] = best_trail.hyperparameters['learning_rate']
        hp_dict['batch_size'] = best_trail.hyperparameters[
            'per_device_train_batch_size']
        hp_dict['n_epoch'] = best_trail.hyperparameters['num_train_epochs']
        hp_dict['seed'] = 42
        with open(hpo_path, 'wb') as fp:
            pickle.dump(hp_dict, fp)
    else:
        logger.info('READING ALREADY SEARCHED HYPERPARAMETERS')
        with open(hpo_path, 'rb') as fp:
            hp_dict = pickle.load(fp)
    return hp_dict


def save_best_checkpoint(checkpoint, save=True):
    if not save:
        logger.info('NOT SAVING TRAINING CHECKPOINT')
        return
    logger.info('SAVING THE BEST CHECKPOINT DURING TRAINING')
    save_checkpoint_path = setting.checkpoint_path
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('prop=%.2f' %
        prop)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('run=%d' % run)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    for filename in os.listdir(checkpoint):
        shutil.move(os.path.join(checkpoint, filename), save_checkpoint_path)


<mask token>

<mask token>


def train_model(trainer):
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and use_checkpoint:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None:
            logger.info(
                f'checkpoint detected, resuming training at {last_checkpoint}')
    if last_checkpoint is not None:
        checkpoint = last_checkpoint
    else:
        checkpoint = None
    trainer.train(resume_from_checkpoint=checkpoint)


<mask token>


def compute_hpo_metric(p):
    label_list = setting.label_list
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    true_predictions = [[label_list[p] for p, l in zip(prediction, label) if
        l != -100] for prediction, label in zip(predictions, labels)]
    true_labels = [[label_list[l] for p, l in zip(prediction, label) if l !=
        -100] for prediction, label in zip(predictions, labels)]
    y_pred_list = list(itertools.chain(*true_predictions))
    y_true_list = list(itertools.chain(*true_labels))
    metric_dict = classification_report(y_true=y_true_list, y_pred=
        y_pred_list, target_names=label_list, output_dict=True)
    weighted_f1 = compute_weighted_f1(metric_dict, target_names=['SN', 'SV'])
    return {'weighted_f1': weighted_f1}


def hpo(tokenized_datasets):
    hpo_path = pathlib.Path('hpo')
    if not hpo_path.exists():
        hpo_path.mkdir()
    hpo_path = hpo_path / pathlib.Path('hpo.pkl')
    if not pathlib.Path(hpo_path).exists():
        logger.info('HYPERPARAMETER SEARCH')
        model_init = lambda : AutoModelForTokenClassification.from_pretrained(
            setting.model_path)
        trainer = Trainer(args=TrainingArguments(output_dir='output/hpo',
            evaluation_strategy='epoch', eval_steps=500, report_to='none',
            disable_tqdm=True), tokenizer=tokenizer, train_dataset=
            tokenized_datasets['train'], eval_dataset=tokenized_datasets[
            'eval'], data_collator=data_collator, model_init=model_init,
            compute_metrics=compute_hpo_metric)
        best_trail = trainer.hyperparameter_search(hp_space=hp_space_optuna,
            direction='maximize', backend='optuna', n_trials=6)
        logger.info('CLEANUP')
        remove_directories(['runs/', 'output/'])
        hp_dict = dict()
        hp_dict['lr'] = best_trail.hyperparameters['learning_rate']
        hp_dict['batch_size'] = best_trail.hyperparameters[
            'per_device_train_batch_size']
        hp_dict['n_epoch'] = best_trail.hyperparameters['num_train_epochs']
        hp_dict['seed'] = 42
        with open(hpo_path, 'wb') as fp:
            pickle.dump(hp_dict, fp)
    else:
        logger.info('READING ALREADY SEARCHED HYPERPARAMETERS')
        with open(hpo_path, 'rb') as fp:
            hp_dict = pickle.load(fp)
    return hp_dict


def save_best_checkpoint(checkpoint, save=True):
    if not save:
        logger.info('NOT SAVING TRAINING CHECKPOINT')
        return
    logger.info('SAVING THE BEST CHECKPOINT DURING TRAINING')
    save_checkpoint_path = setting.checkpoint_path
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('prop=%.2f' %
        prop)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('run=%d' % run)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    for filename in os.listdir(checkpoint):
        shutil.move(os.path.join(checkpoint, filename), save_checkpoint_path)


<mask token>

<mask token>


def initialize_trainer(model, tokenized_datasets, data_collator):
    trainer = Trainer(model=model, args=training_args, tokenizer=tokenizer,
        train_dataset=tokenized_datasets['train'], eval_dataset=
        tokenized_datasets['eval'], data_collator=data_collator,
        compute_metrics=compute_metrics)
    return trainer


def train_model(trainer):
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and use_checkpoint:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None:
            logger.info(
                f'checkpoint detected, resuming training at {last_checkpoint}')
    if last_checkpoint is not None:
        checkpoint = last_checkpoint
    else:
        checkpoint = None
    trainer.train(resume_from_checkpoint=checkpoint)


def get_checkpoint(folder, tokenized_datasets, mode='best', pattern=None):
    assert pattern == None or isinstance(pattern, re.Pattern
        ), 'pattern should be either None or re.Pattern'
    assert mode in ['best', 'median', 'mean', 'worst'
        ], "only 'best', 'median', 'mean', and 'worst' modes are supported"
    checkpoint_name_pattern = re.compile('^' + 'checkpoint' + '\\-(\\d+)$'
        ) if pattern == None else pattern
    checkpoints = [os.path.join(folder, path) for path in os.listdir(folder
        ) if checkpoint_name_pattern.search(path) is not None and os.path.
        isdir(os.path.join(folder, path))]
    checkpoint_dict = dict()
    for checkpoint in checkpoints:
        logger.info('evaluating checkpoint: %s...' % checkpoint)
        model = AutoModelForTokenClassification.from_pretrained(checkpoint)
        trainer = initialize_trainer(model, tokenized_datasets, data_collator)
        metric, metric_str = evaluate_model(trainer, tokenized_datasets,
            name='eval')
        weighted_f1 = compute_weighted_f1(metric, target_names=['SN', 'SV'])
        checkpoint_dict[checkpoint] = weighted_f1
        perf_str = '\n'.join('\t%s: %.4f' % (key, val) for key, val in
            sorted(checkpoint_dict.items(), key=lambda x: x[1], reverse=True))
        logger.info(perf_str)
    if mode == 'best':
        checkpoint = max(checkpoint_dict, key=checkpoint_dict.get)
    else:
        weighted_f1_arr = np.fromiter(checkpoint_dict.values(), dtype=float,
            count=len(checkpoint_dict))
        if mode == 'mean':
            mean = np.mean(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - mean) for key, val in
                checkpoint_dict.items()}
        if mode == 'median':
            median = np.median(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - median) for key, val in
                checkpoint_dict.items()}
        checkpoint = min(checkpoint_dict, key=checkpoint_dict.get)
    return checkpoint


<mask token>


def compute_hpo_metric(p):
    label_list = setting.label_list
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    true_predictions = [[label_list[p] for p, l in zip(prediction, label) if
        l != -100] for prediction, label in zip(predictions, labels)]
    true_labels = [[label_list[l] for p, l in zip(prediction, label) if l !=
        -100] for prediction, label in zip(predictions, labels)]
    y_pred_list = list(itertools.chain(*true_predictions))
    y_true_list = list(itertools.chain(*true_labels))
    metric_dict = classification_report(y_true=y_true_list, y_pred=
        y_pred_list, target_names=label_list, output_dict=True)
    weighted_f1 = compute_weighted_f1(metric_dict, target_names=['SN', 'SV'])
    return {'weighted_f1': weighted_f1}


def hpo(tokenized_datasets):
    hpo_path = pathlib.Path('hpo')
    if not hpo_path.exists():
        hpo_path.mkdir()
    hpo_path = hpo_path / pathlib.Path('hpo.pkl')
    if not pathlib.Path(hpo_path).exists():
        logger.info('HYPERPARAMETER SEARCH')
        model_init = lambda : AutoModelForTokenClassification.from_pretrained(
            setting.model_path)
        trainer = Trainer(args=TrainingArguments(output_dir='output/hpo',
            evaluation_strategy='epoch', eval_steps=500, report_to='none',
            disable_tqdm=True), tokenizer=tokenizer, train_dataset=
            tokenized_datasets['train'], eval_dataset=tokenized_datasets[
            'eval'], data_collator=data_collator, model_init=model_init,
            compute_metrics=compute_hpo_metric)
        best_trail = trainer.hyperparameter_search(hp_space=hp_space_optuna,
            direction='maximize', backend='optuna', n_trials=6)
        logger.info('CLEANUP')
        remove_directories(['runs/', 'output/'])
        hp_dict = dict()
        hp_dict['lr'] = best_trail.hyperparameters['learning_rate']
        hp_dict['batch_size'] = best_trail.hyperparameters[
            'per_device_train_batch_size']
        hp_dict['n_epoch'] = best_trail.hyperparameters['num_train_epochs']
        hp_dict['seed'] = 42
        with open(hpo_path, 'wb') as fp:
            pickle.dump(hp_dict, fp)
    else:
        logger.info('READING ALREADY SEARCHED HYPERPARAMETERS')
        with open(hpo_path, 'rb') as fp:
            hp_dict = pickle.load(fp)
    return hp_dict


def save_best_checkpoint(checkpoint, save=True):
    if not save:
        logger.info('NOT SAVING TRAINING CHECKPOINT')
        return
    logger.info('SAVING THE BEST CHECKPOINT DURING TRAINING')
    save_checkpoint_path = setting.checkpoint_path
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('prop=%.2f' %
        prop)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('run=%d' % run)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    for filename in os.listdir(checkpoint):
        shutil.move(os.path.join(checkpoint, filename), save_checkpoint_path)


<mask token>

<mask token>
assert prop or size, 'prop and size could not be set to None simultaneously'
<mask token>
remove_directories(['data/', 'output/', 'runs/'])
<mask token>
if not log_path.exists():
    log_path.mkdir()
<mask token>
logger.info('\n'.join(settings))
<mask token>


def initialize_trainer(model, tokenized_datasets, data_collator):
    trainer = Trainer(model=model, args=training_args, tokenizer=tokenizer,
        train_dataset=tokenized_datasets['train'], eval_dataset=
        tokenized_datasets['eval'], data_collator=data_collator,
        compute_metrics=compute_metrics)
    return trainer


def train_model(trainer):
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and use_checkpoint:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None:
            logger.info(
                f'checkpoint detected, resuming training at {last_checkpoint}')
    if last_checkpoint is not None:
        checkpoint = last_checkpoint
    else:
        checkpoint = None
    trainer.train(resume_from_checkpoint=checkpoint)


def get_checkpoint(folder, tokenized_datasets, mode='best', pattern=None):
    assert pattern == None or isinstance(pattern, re.Pattern
        ), 'pattern should be either None or re.Pattern'
    assert mode in ['best', 'median', 'mean', 'worst'
        ], "only 'best', 'median', 'mean', and 'worst' modes are supported"
    checkpoint_name_pattern = re.compile('^' + 'checkpoint' + '\\-(\\d+)$'
        ) if pattern == None else pattern
    checkpoints = [os.path.join(folder, path) for path in os.listdir(folder
        ) if checkpoint_name_pattern.search(path) is not None and os.path.
        isdir(os.path.join(folder, path))]
    checkpoint_dict = dict()
    for checkpoint in checkpoints:
        logger.info('evaluating checkpoint: %s...' % checkpoint)
        model = AutoModelForTokenClassification.from_pretrained(checkpoint)
        trainer = initialize_trainer(model, tokenized_datasets, data_collator)
        metric, metric_str = evaluate_model(trainer, tokenized_datasets,
            name='eval')
        weighted_f1 = compute_weighted_f1(metric, target_names=['SN', 'SV'])
        checkpoint_dict[checkpoint] = weighted_f1
        perf_str = '\n'.join('\t%s: %.4f' % (key, val) for key, val in
            sorted(checkpoint_dict.items(), key=lambda x: x[1], reverse=True))
        logger.info(perf_str)
    if mode == 'best':
        checkpoint = max(checkpoint_dict, key=checkpoint_dict.get)
    else:
        weighted_f1_arr = np.fromiter(checkpoint_dict.values(), dtype=float,
            count=len(checkpoint_dict))
        if mode == 'mean':
            mean = np.mean(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - mean) for key, val in
                checkpoint_dict.items()}
        if mode == 'median':
            median = np.median(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - median) for key, val in
                checkpoint_dict.items()}
        checkpoint = min(checkpoint_dict, key=checkpoint_dict.get)
    return checkpoint


def hp_space_optuna(trial):
    return {'learning_rate': trial.suggest_categorical('learning_rate', [
        1e-06, 5e-06, 1e-05]), 'num_train_epochs': trial.
        suggest_categorical('num_train_epochs', [3, 5]),
        'per_device_train_batch_size': trial.suggest_categorical(
        'per_device_train_batch_size', [2])}


def compute_hpo_metric(p):
    label_list = setting.label_list
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    true_predictions = [[label_list[p] for p, l in zip(prediction, label) if
        l != -100] for prediction, label in zip(predictions, labels)]
    true_labels = [[label_list[l] for p, l in zip(prediction, label) if l !=
        -100] for prediction, label in zip(predictions, labels)]
    y_pred_list = list(itertools.chain(*true_predictions))
    y_true_list = list(itertools.chain(*true_labels))
    metric_dict = classification_report(y_true=y_true_list, y_pred=
        y_pred_list, target_names=label_list, output_dict=True)
    weighted_f1 = compute_weighted_f1(metric_dict, target_names=['SN', 'SV'])
    return {'weighted_f1': weighted_f1}


def hpo(tokenized_datasets):
    hpo_path = pathlib.Path('hpo')
    if not hpo_path.exists():
        hpo_path.mkdir()
    hpo_path = hpo_path / pathlib.Path('hpo.pkl')
    if not pathlib.Path(hpo_path).exists():
        logger.info('HYPERPARAMETER SEARCH')
        model_init = lambda : AutoModelForTokenClassification.from_pretrained(
            setting.model_path)
        trainer = Trainer(args=TrainingArguments(output_dir='output/hpo',
            evaluation_strategy='epoch', eval_steps=500, report_to='none',
            disable_tqdm=True), tokenizer=tokenizer, train_dataset=
            tokenized_datasets['train'], eval_dataset=tokenized_datasets[
            'eval'], data_collator=data_collator, model_init=model_init,
            compute_metrics=compute_hpo_metric)
        best_trail = trainer.hyperparameter_search(hp_space=hp_space_optuna,
            direction='maximize', backend='optuna', n_trials=6)
        logger.info('CLEANUP')
        remove_directories(['runs/', 'output/'])
        hp_dict = dict()
        hp_dict['lr'] = best_trail.hyperparameters['learning_rate']
        hp_dict['batch_size'] = best_trail.hyperparameters[
            'per_device_train_batch_size']
        hp_dict['n_epoch'] = best_trail.hyperparameters['num_train_epochs']
        hp_dict['seed'] = 42
        with open(hpo_path, 'wb') as fp:
            pickle.dump(hp_dict, fp)
    else:
        logger.info('READING ALREADY SEARCHED HYPERPARAMETERS')
        with open(hpo_path, 'rb') as fp:
            hp_dict = pickle.load(fp)
    return hp_dict


def save_best_checkpoint(checkpoint, save=True):
    if not save:
        logger.info('NOT SAVING TRAINING CHECKPOINT')
        return
    logger.info('SAVING THE BEST CHECKPOINT DURING TRAINING')
    save_checkpoint_path = setting.checkpoint_path
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('prop=%.2f' %
        prop)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    save_checkpoint_path = save_checkpoint_path / pathlib.Path('run=%d' % run)
    if not save_checkpoint_path.exists():
        save_checkpoint_path.mkdir()
    for filename in os.listdir(checkpoint):
        shutil.move(os.path.join(checkpoint, filename), save_checkpoint_path)


logger.info(f'EXPERIMENT: few-shot learning on {category}')
<mask token>
logger.info('DATASET STATISTICS')
logger.info(dataset.df.name.value_counts())
logger.info(dataset.df[dataset.df.name == 'train'].not_o_only.value_counts(
    normalize=True))
<mask token>
logger.info(
    f"""BEST HYPERPARAMETER: 
	lr={lr:.8f}
	batch_size={batch_size}
	n_epoch={n_epoch}
	seed={seed}"""
    )
<mask token>
wandb.init(project='pretraining', group=
    f'{model_name}_{category}_size-{prop}_frac-{info_frac}', name=
    f'run-{run}', config=param)
<mask token>
logger.info('\n'.join(params))
<mask token>
logger.info('TRAINING: fine tuning on %s' % category)
train_model(trainer)
logger.info('VALIDATION: validate checkponints generated on %s' % category)
<mask token>
logger.info('BEST CHECKPOINT: %s' % best_checkpoint)
logger.info('TESTING')
<mask token>
logger.info(metric_str)
logger.info('METRIC: weighted F1=%.4f' % weighted_f1)
save_best_checkpoint(best_checkpoint, save=True)

import re
import os
import torch
import wandb
import pickle
import shutil
import logging
import pathlib
import itertools

import numpy as np

from datasets import load_dataset
from transformers.trainer_utils import get_last_checkpoint
from transformers import AutoConfig, AutoTokenizer, AutoModelForTokenClassification
from transformers import DataCollatorForTokenClassification, Trainer, TrainingArguments

from sklearn.metrics import classification_report

from utils.data import SimpleNERDataset
from utils.bert import create_tokenized_datasets, evaluate_model, compute_metrics
from utils.common import prepare_logger, remove_directories, compute_weighted_f1

from setting import setting

model_name = setting.model_name
model_name = model_name.replace("/", "-")

# variables controlled by external scripts
# str
category = "memc" if os.environ.get("CATEGORY") is None else os.environ.get("CATEGORY")
# float
prop = 0.10 if os.environ.get("PROP") is None else float(os.environ.get("PROP"))
info_frac = None if (os.environ.get("INFO_FRAC") is None) or (os.environ.get("INFO_FRAC") == str(None)) \
                 else float(os.environ.get("INFO_FRAC")) 
# int
size = None if os.environ.get("SIZE") is None else int(os.environ.get("SIZE"))
run = 1 if os.environ.get("RUN") is None else int(os.environ.get("RUN"))
# boolean
use_checkpoint = True
train_valid_same_size = True if os.environ.get("TRAIN_VALID_SAME_SIZE") == "True" else False

assert prop or size, "prop and size could not be set to None simultaneously"

# variables referenced in this script
n_epoch = 3
n_eval = 10
n_logging = 50
n_checkpoint = 3
n_device = torch.cuda.device_count()

# directories setting
# 1. remove previous experiment artifacts
remove_directories(["data/", "output/", "runs/"])

# 2. make sure all directories exist
log_path = pathlib.Path("logs/")
base_path = pathlib.Path("data/")
if not log_path.exists(): log_path.mkdir()

settings = ["mode=pretraining",
            "category=%s" % category,
            "model=%s" % model_name,
            "prop=%.2f" % prop if prop else "prop=None", 
            "size=%s" % size,
            "info_frac=%.2f" % info_frac if info_frac is not None else "info_frac=None",
            "same_size=%s" % train_valid_same_size,
            "run=%s" % run]

# logging setting
logger = prepare_logger(filename="%s/logfile@%s.log" % (log_path, "_".join(settings)), name="NER")
logger.info("\n".join(settings))

# variables related to transformers
config = AutoConfig.from_pretrained(setting.config_path)
tokenizer = AutoTokenizer.from_pretrained(setting.tokenizer_path, config=config, add_prefix_space=True)
data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=None)

training_args = TrainingArguments(do_train=True,
                                  do_eval=True,
                                  num_train_epochs=n_epoch,
                                  output_dir="output/ner",
                                  per_device_train_batch_size=2,
                                  per_device_eval_batch_size=2,
                                  # fp16
                                  fp16=True,
                                  # logging
                                  logging_strategy="steps",
                                  logging_steps=10,
                                  disable_tqdm=True,
                                  # checkpoint setting
                                  save_strategy="steps",
                                  save_steps=100,
                                  # evaluation
                                  evaluation_strategy="steps",
                                  eval_steps=100,
                                  # wandb
                                  report_to="wandb")


def initialize_trainer(model, tokenized_datasets, data_collator):
    trainer = Trainer(model=model, 
                      args=training_args,
                      tokenizer=tokenizer,
                      train_dataset=tokenized_datasets["train"],
                      eval_dataset=tokenized_datasets["eval"],
                      data_collator=data_collator,
                      compute_metrics=compute_metrics)
    return trainer


def train_model(trainer):
    # load checkpoint and train
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and use_checkpoint:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None:
            logger.info(f"checkpoint detected, resuming training at {last_checkpoint}")

    # if use_checkpoint == False, then checkpoint == None, no checkpoints will be loaded
    if last_checkpoint is not None:
        checkpoint = last_checkpoint
    else:
        checkpoint = None
        
    trainer.train(resume_from_checkpoint=checkpoint)


def get_checkpoint(folder, tokenized_datasets, mode="best", pattern=None):
    # checkpoints: a list of directories
    # for pretrained checkpoints, pattern: re.compile(r"run=(\d+)")
    assert (pattern == None) or isinstance(pattern, re.Pattern), "pattern should be either None or re.Pattern"
    assert mode in ["best", "median", "mean", "worst"], "only 'best', 'median', 'mean', and 'worst' modes are supported"

    checkpoint_name_pattern = re.compile(r"^" + "checkpoint" + r"\-(\d+)$") if pattern == None else pattern
    checkpoints = [os.path.join(folder, path) for path in os.listdir(folder) if (checkpoint_name_pattern.search(path) is not None) and 
                                                                                os.path.isdir(os.path.join(folder, path))]
    checkpoint_dict = dict()
    for checkpoint in checkpoints:
        logger.info("evaluating checkpoint: %s..." % checkpoint)
        model = AutoModelForTokenClassification.from_pretrained(checkpoint)

        # NOTE: as initialize_trainer() is used here, could NOT put this function into utils.bert
        trainer = initialize_trainer(model, tokenized_datasets, data_collator)
        metric, metric_str = evaluate_model(trainer, tokenized_datasets, name="eval")

        weighted_f1 = compute_weighted_f1(metric, target_names=["SN", "SV"])

        checkpoint_dict[checkpoint] = weighted_f1
        # print the current evaluated checkpoints based on descending order of weighted_f1
        perf_str = "\n".join("\t%s: %.4f" % (key, val) 
                             for key, val in sorted(checkpoint_dict.items(), key=lambda x: x[1], reverse=True))
        logger.info(perf_str)


    # select checkpoints based on different criterion
    if mode == "best":
        checkpoint = max(checkpoint_dict, key=checkpoint_dict.get)
    else:
        weighted_f1_arr = np.fromiter(checkpoint_dict.values(), dtype=float, count=len(checkpoint_dict))
        if mode == "mean":
            mean = np.mean(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - mean) for key, val in checkpoint_dict.items()}
        if mode == "median":
            median = np.median(weighted_f1_arr)
            checkpoint_dict = {key: abs(val - median) for key, val in checkpoint_dict.items()}
        # if mode == "worst", no need to modify the checkpoint_dict
        checkpoint = min(checkpoint_dict, key=checkpoint_dict.get)
    
    return checkpoint


def hp_space_optuna(trial):
    return {
        "learning_rate": trial.suggest_categorical("learning_rate", [1e-6, 5e-6, 1e-5]),
        "num_train_epochs": trial.suggest_categorical("num_train_epochs", [3, 5]),
        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [2]),
    }


def compute_hpo_metric(p):
    label_list = setting.label_list

    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    y_pred_list = list(itertools.chain(*true_predictions))
    y_true_list = list(itertools.chain(*true_labels))

    metric_dict = classification_report(y_true=y_true_list, 
                                        y_pred=y_pred_list, 
                                        target_names=label_list, 
                                        output_dict=True)

    weighted_f1 = compute_weighted_f1(metric_dict, target_names=["SN", "SV"])

    return {"weighted_f1": weighted_f1}


def hpo(tokenized_datasets):
    hpo_path = pathlib.Path("hpo")
    if not hpo_path.exists(): hpo_path.mkdir()
    hpo_path = hpo_path / pathlib.Path("hpo.pkl")

    if not pathlib.Path(hpo_path).exists(): 
        logger.info("HYPERPARAMETER SEARCH")
        model_init = lambda: AutoModelForTokenClassification.from_pretrained(setting.model_path)
        trainer = Trainer(args=TrainingArguments(output_dir="output/hpo", evaluation_strategy="epoch", eval_steps=500, report_to="none", disable_tqdm=True),
                          tokenizer=tokenizer,
                          train_dataset=tokenized_datasets["train"],
                          eval_dataset=tokenized_datasets["eval"],
                          data_collator=data_collator,
                          model_init=model_init,
                          compute_metrics=compute_hpo_metric)

        best_trail = trainer.hyperparameter_search(hp_space=hp_space_optuna,
                                                   direction="maximize", 
                                                   backend="optuna", 
                                                   n_trials=6)

        logger.info("CLEANUP")
        remove_directories(["runs/", "output/"])

        hp_dict = dict()

        hp_dict["lr"] = best_trail.hyperparameters["learning_rate"]
        hp_dict["batch_size"] = best_trail.hyperparameters["per_device_train_batch_size"]
        hp_dict["n_epoch"] = best_trail.hyperparameters["num_train_epochs"]
        hp_dict["seed"] = 42

        with open(hpo_path, "wb") as fp:
            pickle.dump(hp_dict, fp)
    else:
        logger.info("READING ALREADY SEARCHED HYPERPARAMETERS")
        with open(hpo_path, "rb") as fp:
            hp_dict = pickle.load(fp)
    
    return hp_dict


def save_best_checkpoint(checkpoint, save=True):
    # checkpoint: the filename of best checkpoint during evaluation
    if not save: 
        logger.info("NOT SAVING TRAINING CHECKPOINT")
        return 

    logger.info("SAVING THE BEST CHECKPOINT DURING TRAINING")
    save_checkpoint_path = setting.checkpoint_path
    if not save_checkpoint_path.exists(): save_checkpoint_path.mkdir()

    save_checkpoint_path = save_checkpoint_path / pathlib.Path("prop=%.2f" % prop)
    if not save_checkpoint_path.exists(): save_checkpoint_path.mkdir()

    save_checkpoint_path = save_checkpoint_path / pathlib.Path("run=%d" % run)
    if not save_checkpoint_path.exists(): save_checkpoint_path.mkdir()

    for filename in os.listdir(checkpoint):
        shutil.move(os.path.join(checkpoint, filename), save_checkpoint_path)
    

logger.info(f"EXPERIMENT: few-shot learning on {category}")
# data
dataset = SimpleNERDataset(category=category, train_size=size, train_frac=prop, 
                           info_frac=info_frac, train_valid_same_size=train_valid_same_size)
tokenized_datasets = create_tokenized_datasets(dataset, in_memory=True)

logger.info("DATASET STATISTICS")
logger.info(dataset.df.name.value_counts())
logger.info(dataset.df[dataset.df.name == "train"].not_o_only.value_counts(normalize=True))

true_info_frac = dataset.df[dataset.df.name == "train"].not_o_only.value_counts(normalize=True)[True]

# search for best main checkpoint and hyperparameters
hp_dict = hpo(tokenized_datasets)

lr = hp_dict["lr"] if hp_dict is not None else 1e-5
batch_size = hp_dict["batch_size"] if hp_dict is not None else 2
n_epoch = hp_dict["n_epoch"] if hp_dict is not None else 3
seed = hp_dict["seed"] if hp_dict is not None else 42

logger.info(f"BEST HYPERPARAMETER: \n\tlr={lr:.8f}\n\tbatch_size={batch_size}\n\tn_epoch={n_epoch}\n\tseed={seed}")

training_args.learning_rate = lr
training_args.num_train_epochs = n_epoch
training_args.per_device_train_batch_size = batch_size
training_args.per_device_eval_batch_size = batch_size
training_args.seed = seed


param = {"lr": lr, 
         "batch_size": batch_size, 
         "n_epoch": n_epoch, 
         "seed": seed, 
         "category": category,
         "train_size": len(dataset.df[dataset.df.name == "train"]),
         "valid_size": len(dataset.df[dataset.df.name == "eval"]),
         "info_frac": true_info_frac,
         "run": run}

# wandb
wandb.init(project="pretraining", 
           group=f"{model_name}_{category}_size-{prop}_frac-{info_frac}", 
           name=f"run-{run}",
           config=param)

# model
model = AutoModelForTokenClassification.from_pretrained(setting.model_path)
trainer = initialize_trainer(model, tokenized_datasets, data_collator)

# training_args
train_size = len(dataset.train_idx)
training_args.save_steps = max([int(train_size * n_epoch / (training_args.per_device_train_batch_size * n_device) / n_checkpoint), 1])
training_args.eval_steps = max([int(train_size * n_epoch / (training_args.per_device_train_batch_size * n_device) / n_eval), 1])
training_args.logging_steps = max([int(train_size * n_epoch / (training_args.per_device_train_batch_size * n_device) / n_logging), 1])

# logging parameters
params = [f"TRAIN SIZE: {train_size}",
          f"BATCH SIZE: {training_args.per_device_eval_batch_size}",
          f"DEVICE NUM: {n_device}"]
logger.info("\n".join(params))

# trainer
trainer = initialize_trainer(model, tokenized_datasets, data_collator)

# train-val-test
logger.info("TRAINING: fine tuning on %s" % category)
train_model(trainer)

logger.info("VALIDATION: validate checkponints generated on %s" % category)
best_checkpoint = get_checkpoint(training_args.output_dir, 
                                 tokenized_datasets,
                                 mode="best",
                                 pattern=None)
logger.info("BEST CHECKPOINT: %s" % best_checkpoint)

logger.info("TESTING")
model = AutoModelForTokenClassification.from_pretrained(best_checkpoint)
trainer = initialize_trainer(model, tokenized_datasets, data_collator)

metric, metric_str = evaluate_model(trainer, tokenized_datasets, name="test")
weighted_f1 = compute_weighted_f1(metric, target_names=["SN", "SV"])
logger.info(metric_str)
logger.info("METRIC: weighted F1=%.4f" % weighted_f1)

save_best_checkpoint(best_checkpoint, save=True)
