main.py

import logging
#import math
import os
import sys
#from itertools import chain
import warnings
import pickle
import io
import numpy as np
import time
import types
import random

import datasets
from datasets import concatenate_datasets, load_dataset, load_from_disk#, load_metric

import transformers
from transformers import (
    #CONFIG_MAPPING,
    #MODEL_FOR_CAUSAL_LM_MAPPING,
    AutoConfig,
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    #TrainingArguments,
    default_data_collator,
    # is_torch_tpu_available,
    set_seed,
)
from transformers.utils import check_min_version

# prompt learning
# from peft import get_peft_config, get_peft_model, PromptTuningInit, PromptTuningConfig, TaskType, PeftType, PeftModelForCausalLM

import torch as tc
from torch.utils.data import DataLoader

import models

# from nltk.corpus import wordnet

from transformers.testing_utils import CaptureLogger
# from transformers.trainer_utils import get_last_checkpoint
# from transformers.utils import check_min_version, send_example_telemetry
# from transformers.utils.versions import require_version

from args import *
import util

from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
from tqdm import tqdm

def init_datasets_qna(data_args, model_args, training_args):  
    # INFO: Deprecated, this is for nli dataset
    if data_args.dataset_config_name:
        raw_datasets = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            # cache_dir=model_args.cache_dir,
            trust_remote_code=True,
        )
        return raw_datasets['validation']
    
    try:
        raw_datasets = load_from_disk(
            data_args.dataset_name,
        )
    except:
        if not training_args.method.startswith('Dialog'):
            raise FileNotFoundError(f"Dataset {data_args.dataset_name} not found in disk.")
        return None
    
    # if not training_args.method.startswith('Dialog'):
    #     raw_datasets = raw_datasets.shuffle(seed=training_args.seed)
    raw_datasets = raw_datasets.shuffle(seed=training_args.seed)
    # FIXME: will be deleted
    if training_args.method == 'Static-Dialog-OSG-Semi-Bandit':
        raw_datasets = raw_datasets.select(range(training_args.static_dialog_chunk, training_args.static_dialog_chunk + 500))
        print(f"Static Dialog dataset chunk: {training_args.static_dialog_chunk}~{training_args.static_dialog_chunk+499}, size: {len(raw_datasets)}")
    return raw_datasets

def build_shifted_dataset(nq_ds, trivia_ds, chunk_size, num_chunks, first_flag, seed):
    chunks = []
    for idx in range(num_chunks):
        base_seed = seed + idx
        
        if chunk_size <= len(nq_ds):
            nq_chunk = nq_ds.shuffle(seed=base_seed).select(range(chunk_size))
        else: 
            reps_nq = (chunk_size + len(nq_ds) - 1) // len(nq_ds)
            nq_pieces = [nq_ds.shuffle(seed=base_seed + r) for r in range(reps_nq)]
            nq_big = concatenate_datasets(nq_pieces)
            nq_chunk = nq_big.select(range(chunk_size))  
        
        if chunk_size <= len(trivia_ds):
            tri_chunk = trivia_ds.shuffle(seed=base_seed).select(range(chunk_size))
        else: 
            reps_tri = (chunk_size + len(trivia_ds) - 1) // len(trivia_ds)
            tri_pieces = [trivia_ds.shuffle(seed=base_seed + r) for r in range(reps_tri)]
            tri_big = concatenate_datasets(tri_pieces)
            tri_chunk = tri_big.select(range(chunk_size)) 
            
        if (first_flag=='tri'):
            chunks.extend([tri_chunk, nq_chunk])
        else: 
            chunks.extend([nq_chunk, tri_chunk])
    return concatenate_datasets(chunks)

#==================================================
# training andevaluation code
#==================================================
def main():       
    # read args
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, UncertaintyTrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # setup logger
    os.makedirs(os.path.join(training_args.snapshot_root, training_args.method, training_args.exp_name), exist_ok=True)
    sys.stdout = util.Logger(os.path.join(training_args.snapshot_root, training_args.method,training_args.exp_name, 'out'))    
    sys.stderr = util.Logger(os.path.join(training_args.snapshot_root, training_args.method,training_args.exp_name, 'out_err'))

    if data_args.log_args is True:
        print('==================================================')
        print(model_args)
        print(training_args)
        if 'data_args' in locals():
            print(data_args)
        
    logger = logging.getLogger(__name__)

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()
    
    # init datasets
    if training_args.method.startswith('Shift'):
        nq_dataset_path = data_args.nq_dataset_name
        triviaqa_dataset_path = data_args.triviaqa_dataset_name
    else:
        raw_dataset = init_datasets_qna(data_args, model_args, training_args)

    np.random.seed(training_args.seed) # for unified restuls, will not affect much.
    
    # FIXME: Deprecated?
    if training_args.method.startswith('IID'):
        # INFO: IID experiment
        if training_args.method == 'IID-OSG':
            l = models.OSGLearner(
                params=training_args,
                name_postfix='osg'
            )
        elif training_args.method == 'IID-OSG-Bandit':
            l = models.OSGAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg'
            )
        elif training_args.method == 'IID-OSG-Semi-Bandit':
            l = models.OSGSemiAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg_bandit'
            )
        else:
            raise NotImplementedError
        
        l.train(
            raw_dataset,
            updated_params=types.SimpleNamespace(n=len(raw_dataset))
        )
    
    elif training_args.method == 'Plot':
        # INFO: Violin plot experiment
        fdrs1, fdrs2, fdrs3 = [], [], []
        
        for idx in tqdm(range(100)):
            base_seed = training_args.seed + idx
            size = len(raw_dataset)
            T = training_args.T

            if T <= size:
                dataset = raw_dataset.shuffle(seed=base_seed).select(range(T))
            else:
                reps = (T + size - 1) // size 
                pieces = [raw_dataset.shuffle(seed=base_seed + r) for r in range(reps)]
                big = concatenate_datasets(pieces)
                dataset = big.select(range(T))            
            # dataset = raw_dataset.shuffle(seed=training_args.seed+idx).select(range(training_args.T))
            # np.random.seed(training_args.seed+idx)
            
            l1, l2, l3 = models.OSGLearner(params=training_args, name_postfix='osg'), models.OSGAdvBanditLearner(params=training_args, name_postfix='adv_sg'), models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            
            results1 = l1.train(dataset, updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            results2 = l2.train(dataset, updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            results3 = l3.train(dataset,updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            
            fdrs1.append(results1['fdrs'][-1])
            fdrs2.append(results2['fdrs'][-1])
            fdrs3.append(results3['fdrs'][-1])
            
        json.dump(fdrs1, open(os.path.join(training_args.output_dir, 'fdrs1.json'), 'w'))
        json.dump(fdrs2, open(os.path.join(training_args.output_dir, 'fdrs2.json'), 'w'))
        json.dump(fdrs3, open(os.path.join(training_args.output_dir, 'fdrs3.json'), 'w'))
             
        with PdfPages(os.path.join(training_args.output_dir, 'violin.pdf')) as pdf:
            os.makedirs(training_args.output_dir, exist_ok=True)
            plt.figure()
            vi = plt.violinplot([fdrs1, fdrs2, fdrs3], showmeans=True, showextrema=False)
            plt.xlabel('Method', fontsize=15)
            plt.xticks(np.arange(1,4), [r'$\mathtt{EW\text{-}SG}$', r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', r'$\mathtt{ExSUL}$ (ours)'])
            plt.ylabel('FDR', fontsize=15)
            
            if training_args.model == 'llama3.1-8B' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.239
            elif training_args.model == 'gpt-3.5' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.124
            elif training_args.model == 'llama3.1-8B' and 'nq' in data_args.dataset_name:
                ylim = 0.425
            elif training_args.model == 'gpt-3.5' and 'nq' in data_args.dataset_name:
                ylim = 0.294
                
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.rcParams.update({'font.size': 14})
            
            vi['bodies'][0].set_facecolor('aqua')
            vi['bodies'][1].set_facecolor('red')
            vi['bodies'][2].set_facecolor('lightgreen')
            
            plt.savefig(os.path.join(training_args.output_dir, 'violin.png'))
            pdf.savefig(bbox_inches='tight')
            
    elif training_args.method == 'Plot-over-alpha-parallel':
        # INFO: Violin plot of fdr, efficiency
            
        # FIXME trivial implementation
        # if training_args.model == 'llama3.1-8B' and 'triviaqa' in data_args.dataset_name:
        #     alphas = [0.1, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16]
        # elif training_args.model == 'gpt-3.5' and 'triviaqa' in data_args.dataset_name:
        #     alphas = [0.25, 0.26, 0.27, 0.28, 0.29, 0.3]
        # elif training_args.model == 'llama3.1-8B' and 'nq' in data_args.dataset_name:
        #     alphas = [0.25, 0.26, 0.27, 0.28, 0.29, 0.3]
        # elif training_args.model == 'gpt-3.5' and 'nq' in data_args.dataset_name:
        #     alphas = [0.25, 0.26, 0.27, 0.28, 0.29, 0.3]
        alphas = training_args.ablation_alphas
        
        # fdr1_total = []
        # fdr2_total = []
        # fdr3_total = []
        # fdr4_total = []
        # ineff1_total = []
        # ineff2_total = []
        # ineff3_total = []
        # ineff4_total = []
        # for aidx, alpha in enumerate(alphas):
        fdrs1 = []
        fdrs2 = []
        fdrs3 = []
        ineff1 = []
        ineff2 = []
        ineff3 = []
        
        if os.path.exists(os.path.join(training_args.output_dir, f'fdr3_{training_args.alpha}.json')):
            return
        
        for idx in tqdm(range(30)):
            dataset = raw_dataset.shuffle(seed=training_args.seed+idx).select(range(training_args.T))
            np.random.seed(training_args.seed+idx)
            
            l3 = models.OSGLearner(
                params=training_args,
                name_postfix='osg'
            )
            l2 = models.OSGAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg'
            )
            l1 = models.OSGSemiAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg_bandit'
            )
            
            results1 = l1.train(
                dataset,
                updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx, alpha=training_args.alpha)
            )
            results2 = l2.train(
                dataset,
                updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx, alpha=training_args.alpha)
            )
            results3 = l3.train(
                dataset,
                updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx, alpha=training_args.alpha)
            )
            
            fdrs1.append(results1['fdrs'][-1])
            fdrs2.append(results2['fdrs'][-1])
            fdrs3.append(results3['fdrs'][-1])
            ineff1.append(1-results1['efficiencies'][-1])
            ineff2.append(1-results2['efficiencies'][-1])
            ineff3.append(1-results3['efficiencies'][-1])
        
            # fdr1_total.append(fdrs1)
            # fdr2_total.append(fdrs2)
            # fdr3_total.append(fdrs3)
            # fdr4_total.append(fdrs4)
            # ineff1_total.append(ineff1)
            # ineff2_total.append(ineff2)
            # ineff3_total.append(ineff3)
            # ineff4_total.append(ineff4)
            
        json.dump(fdrs1, open(os.path.join(training_args.output_dir, f'fdr1_{training_args.alpha}.json'), 'w'))
        json.dump(fdrs2, open(os.path.join(training_args.output_dir, f'fdr2_{training_args.alpha}.json'), 'w'))
        json.dump(fdrs3, open(os.path.join(training_args.output_dir, f'fdr3_{training_args.alpha}.json'), 'w'))
        json.dump(ineff1, open(os.path.join(training_args.output_dir, f'ineff1_{training_args.alpha}.json'), 'w'))
        json.dump(ineff2, open(os.path.join(training_args.output_dir, f'ineff2_{training_args.alpha}.json'), 'w'))
        json.dump(ineff3, open(os.path.join(training_args.output_dir, f'ineff3_{training_args.alpha}.json'), 'w'))
        
        # with PdfPages(os.path.join(training_args.output_dir, 'violin_fdr.pdf')) as pdf:
        #     plt.figure()
        #     plt.xlabel(r'$\alpha$', fontsize=15)
        #     plt.xticks(np.arange(1,7), alphas)
        #     plt.ylabel('FDR', fontsize=15)
        #     # plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

        #     # plt.title('FDR')
        #     plt.legend(loc='upper right', fontsize=10)
        #     plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
        #     plt.tight_layout()
        #     # plt.rcParams.update({'font.size': 14})
        #     vi = plt.violinplot(fdr3_total, showmeans=True, showextrema=False)
                
        #     plt.plot([0.2, 0.35], [0.2, 0.35], linestyle='--', color='k', linewidth=0.5, label=r'$y=x$')
        #     # vi['bodies'][0].set_facecolor('aqua')
        #     # vi['bodies'][1].set_facecolor('red')
        #     [d.set_facecolor('lightgreen') for d in vi['bodies']]
    
        #     plt.savefig(os.path.join(training_args.output_dir, 'violin_fdr.png'))
        #     pdf.savefig(bbox_inches='tight')
        
        # with PdfPages(os.path.join(training_args.output_dir, 'violin_ineff.pdf')) as pdf:
        #     plt.figure()
        #     plt.xlabel(r'$\alpha$', fontsize=15)
        #     plt.xticks(np.arange(1,7), alphas)
        #     plt.ylabel('Inefficiency', fontsize=15)
        #     # plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

        #     # plt.title('FDR')
        #     plt.legend(loc='upper right', fontsize=10)
        #     plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
        #     plt.tight_layout()
        #     # plt.rcParams.update({'font.size': 14})
        #     vi = plt.violinplot(ineff3_total, showmeans=True, showextrema=False)
                
        #     # vi['bodies'][0].set_facecolor('aqua')
        #     # vi['bodies'][1].set_facecolor('red')
        #     [d.set_facecolor('lightgreen') for d in vi['bodies']]
    
        #     plt.savefig(os.path.join(training_args.output_dir, 'violin_ineff.png'))
        #     pdf.savefig(bbox_inches='tight')
            
    elif training_args.method == 'Plot-over-alpha':
        # INFO: Violin plot of fdr, efficiency
            
        # FIXME trivial implementation
        # if training_args.model == 'llama3.1-8B' and 'triviaqa' in data_args.dataset_name:
        #     alphas = [0.1, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16]
        # elif training_args.model == 'gpt-3.5' and 'triviaqa' in data_args.dataset_name:
        #     alphas = [0.25, 0.26, 0.27, 0.28, 0.29, 0.3]
        # elif training_args.model == 'llama3.1-8B' and 'nq' in data_args.dataset_name:
        #     alphas = [0.25, 0.26, 0.27, 0.28, 0.29, 0.3]
        # elif training_args.model == 'gpt-3.5' and 'nq' in data_args.dataset_name:
        #     alphas = [0.25, 0.26, 0.27, 0.28, 0.29, 0.3]
        alphas = training_args.ablation_alphas
        
        fdr1_total = []
        fdr2_total = []
        fdr3_total = []
        fdr4_total = []
        ineff1_total = []
        ineff2_total = []
        ineff3_total = []
        ineff4_total = []
        for aidx, alpha in enumerate(alphas):
            fdrs1 = []
            fdrs2 = []
            fdrs3 = []
            fdrs4 = []
            ineff1 = []
            ineff2 = []
            ineff3 = []
            ineff4 = []
            for idx in tqdm(range(30)):
                dataset = raw_dataset.shuffle(seed=training_args.seed+idx+100*aidx).select(range(training_args.T))
                np.random.seed(training_args.seed+idx+100*aidx)
                
                l1 = models.OSGLearner(
                    params=training_args,
                    name_postfix='osg'
                )
                l2 = models.OSGAdvBanditLearner(
                    params=training_args,
                    name_postfix='adv_sg'
                )
                l3 = models.OSGSemiAdvBanditLearner(
                    params=training_args,
                    name_postfix='adv_sg_bandit'
                )
                l4 = models.Exp3EwLearner(
                    params=training_args,
                    name_postfix='adv_sg_bandit_expew'
                )
                
                results1 = l1.train(
                    dataset,
                    updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx+100*aidx, alpha=alpha)
                )
                results2 = l2.train(
                    dataset,
                    updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx+100*aidx, alpha=alpha)
                )
                results3 = l3.train(
                    dataset,
                    updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx+100*aidx, alpha=alpha)
                )
                results4 = l4.train(
                    dataset,
                    updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx+100*aidx, alpha=alpha)
                )
                
                fdrs1.append(results1['fdrs'][-1])
                fdrs2.append(results2['fdrs'][-1])
                fdrs3.append(results3['fdrs'][-1])
                fdrs4.append(results4['fdrs'][-1])
                ineff1.append(1-results1['efficiencies'][-1])
                ineff2.append(1-results2['efficiencies'][-1])
                ineff3.append(1-results3['efficiencies'][-1])
                ineff4.append(1-results4['efficiencies'][-1])
            
            fdr1_total.append(fdrs1)
            fdr2_total.append(fdrs2)
            fdr3_total.append(fdrs3)
            fdr4_total.append(fdrs4)
            ineff1_total.append(ineff1)
            ineff2_total.append(ineff2)
            ineff3_total.append(ineff3)
            ineff4_total.append(ineff4)
            
        json.dump(fdr1_total, open(os.path.join(training_args.output_dir, 'fdr1_total.json'), 'w'))
        json.dump(fdr2_total, open(os.path.join(training_args.output_dir, 'fdr2_total.json'), 'w'))
        json.dump(fdr3_total, open(os.path.join(training_args.output_dir, 'fdr3_total.json'), 'w'))
        json.dump(fdr4_total, open(os.path.join(training_args.output_dir, 'fdr4_total.json'), 'w'))
        json.dump(ineff1_total, open(os.path.join(training_args.output_dir, 'ineff1_total.json'), 'w'))
        json.dump(ineff2_total, open(os.path.join(training_args.output_dir, 'ineff2_total.json'), 'w'))
        json.dump(ineff3_total, open(os.path.join(training_args.output_dir, 'ineff3_total.json'), 'w'))
        json.dump(ineff4_total, open(os.path.join(training_args.output_dir, 'ineff4_total.json'), 'w'))
        
        with PdfPages(os.path.join(training_args.output_dir, 'violin_fdr.pdf')) as pdf:
            plt.figure()
            plt.xlabel(r'$\alpha$', fontsize=15)
            plt.xticks(np.arange(1,7), alphas)
            plt.ylabel('FDR', fontsize=15)
            # plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            vi = plt.violinplot(fdr3_total, showmeans=True, showextrema=False)
                
            plt.plot([0.2, 0.35], [0.2, 0.35], linestyle='--', color='k', linewidth=0.5, label=r'$y=x$')
            # vi['bodies'][0].set_facecolor('aqua')
            # vi['bodies'][1].set_facecolor('red')
            [d.set_facecolor('lightgreen') for d in vi['bodies']]
    
            plt.savefig(os.path.join(training_args.output_dir, 'violin_fdr.png'))
            pdf.savefig(bbox_inches='tight')
        
        with PdfPages(os.path.join(training_args.output_dir, 'violin_ineff.pdf')) as pdf:
            plt.figure()
            plt.xlabel(r'$\alpha$', fontsize=15)
            plt.xticks(np.arange(1,7), alphas)
            plt.ylabel('Inefficiency', fontsize=15)
            # plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            vi = plt.violinplot(ineff3_total, showmeans=True, showextrema=False)
                
            # vi['bodies'][0].set_facecolor('aqua')
            # vi['bodies'][1].set_facecolor('red')
            [d.set_facecolor('lightgreen') for d in vi['bodies']]
    
            plt.savefig(os.path.join(training_args.output_dir, 'violin_ineff.png'))
            pdf.savefig(bbox_inches='tight')
        
        with PdfPages(os.path.join(training_args.output_dir, 'violin_fdr_expew.pdf')) as pdf:
            plt.figure()
            plt.xlabel(r'$\alpha$', fontsize=15)
            plt.xticks(np.arange(1,7), alphas)
            plt.ylabel('FDR', fontsize=15)
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            vi = plt.violinplot(fdr4_total, showmeans=True, showextrema=False)
                
                
            # vi['bodies'][0].set_facecolor('aqua')
            # vi['bodies'][1].set_facecolor('red')
            [d.set_facecolor('lightgreen') for d in vi['bodies']]
    
            plt.savefig(os.path.join(training_args.output_dir, 'violin_fdr_expew.png'))
            pdf.savefig(bbox_inches='tight')
        
        with PdfPages(os.path.join(training_args.output_dir, 'violin_ineff_expew.pdf')) as pdf:
            plt.figure()
            plt.xlabel(r'$\alpha$', fontsize=15)
            plt.xticks(np.arange(1,7), alphas)
            plt.ylabel('Inefficiency', fontsize=15)
            # plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')

            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            vi = plt.violinplot(ineff4_total, showmeans=True, showextrema=False)
                
            # vi['bodies'][0].set_facecolor('aqua')
            # vi['bodies'][1].set_facecolor('red')
            [d.set_facecolor('lightgreen') for d in vi['bodies']]
    
            plt.savefig(os.path.join(training_args.output_dir, 'violin_ineff_expew.png'))
            pdf.savefig(bbox_inches='tight')
                 
    elif training_args.method == 'All':
        l1, l2, l3 = models.OSGLearner(params=training_args, name_postfix='osg'), models.OSGAdvBanditLearner(params=training_args, name_postfix='adv_sg'), models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
        
        results1, results2, results3 = l1.train(raw_dataset, updated_params=types.SimpleNamespace(n=len(raw_dataset), seed=training_args.seed)), l2.train(raw_dataset, updated_params=types.SimpleNamespace(n=len(raw_dataset), seed=training_args.seed)), l3.train(raw_dataset, updated_params=types.SimpleNamespace(n=len(raw_dataset), seed=training_args.seed))
        
        # TODO: Reduce lines
        with PdfPages(os.path.join(training_args.output_dir, 'fdrs_plot.pdf')) as pdf:
            plt.figure()
            plt.plot(results3['fdrs'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(results2['fdrs'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(results1['fdrs'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            
            assert results1['empirical_fdrs'] == results2['empirical_fdrs'] == results3['empirical_fdrs'], "Empirical FDRs should be the same across methods."
            plt.plot(results3['empirical_fdrs'], label=r'$\mathtt{No\text{-}SG}$', color='orange')
            
            if training_args.model == 'llama3.1-8B' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.239
            elif training_args.model == 'gpt-3.5' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.124
            elif training_args.model == 'llama3.1-8B' and 'nq' in data_args.dataset_name:
                ylim = 0.425
            elif training_args.model == 'gpt-3.5' and 'nq' in data_args.dataset_name:
                ylim = 0.294
            elif training_args.model == 'gpt-3.5' and 'static' in data_args.dataset_name:
                ylim = 0.294
            # ylim = 0.6
                
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('FDR', fontsize=15)
            # plt.title('FDR')
            plt.ylim(training_args.alpha - (ylim-training_args.alpha)*0.9, ylim + (ylim-training_args.alpha)*0.9)
            # plt.ylim(0, 0.2) if 'triviaqa' in data_args.dataset_name else plt.ylim(0, 0.5)
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
            # plt.axhline(y=ylim, color='k', alpha=0.7, linestyle='dotted', label='Empirical FDR upper-bound')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'fdrs_plot.png'))
            pdf.savefig(bbox_inches='tight')
        
        with PdfPages(os.path.join(training_args.output_dir, 'inefficiencies_plot.pdf')) as pdf:
            plt.figure()
            plt.plot(1-np.array(results3['efficiencies']), label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(1-np.array(results2['efficiencies']), label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(1-np.array(results1['efficiencies']), label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('Inefficiency', fontsize=15)
            # plt.axhline(y=results1['best_efficiency'], color='k', linestyle='dashed', label='Best Efficiency')
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'inefficiencies_plot.png'))
            pdf.savefig(bbox_inches='tight')
        
        
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_osg.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['weights'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_osg.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_adv.pdf')) as pdf:
            plt.figure()
            plt.plot(results2['weights'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_adv.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_semi.pdf')) as pdf:
            plt.figure()
            plt.plot(results3['weights'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_semi.png'))
            pdf.savefig(bbox_inches='tight')    
        
    elif training_args.method == 'Ablation':
        l1, l2= models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit'), models.OSGAblationLearner(params=training_args, name_postfix='adv_sg_bandit_ablation')
        
        results1, results2 = l1.train(raw_dataset, updated_params=types.SimpleNamespace(n=len(raw_dataset), seed=training_args.seed)), l2.train(raw_dataset, updated_params=types.SimpleNamespace(n=len(raw_dataset), seed=training_args.seed))
        
        # TODO: Reduce lines
        with PdfPages(os.path.join(training_args.output_dir, 'fdrs_plot.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['fdrs'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(results2['fdrs'], label=r'$\mathtt{ExSUL}$ w/o $a_t$', color='purple')
            assert results1['empirical_fdrs'] == results2['empirical_fdrs'], "Empirical FDRs should be the same across methods."

            if training_args.model == 'llama3.1-8B' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.239
            elif training_args.model == 'gpt-3.5' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.124
            elif training_args.model == 'llama3.1-8B' and 'nq' in data_args.dataset_name:
                ylim = 0.425
            elif training_args.model == 'gpt-3.5' and 'nq' in data_args.dataset_name:
                ylim = 0.294
            elif training_args.model == 'gpt-3.5' and 'static' in data_args.dataset_name:
                ylim = 0.294
                
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('FDR', fontsize=15)
            # plt.title('FDR')
            plt.ylim(training_args.alpha - (ylim-training_args.alpha)*0.9, ylim + (ylim-training_args.alpha)*0.9)
            # plt.ylim(0, 0.2) if 'triviaqa' in data_args.dataset_name else plt.ylim(0, 0.5)
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
            # plt.axhline(y=ylim, color='k', alpha=0.7, linestyle='dotted', label='Empirical FDR upper-bound')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'fdrs_plot.png'))
            pdf.savefig(bbox_inches='tight')
        
        with PdfPages(os.path.join(training_args.output_dir, 'inefficiencies_plot.pdf')) as pdf:
            plt.figure()
            plt.plot(1-np.array(results1['efficiencies']), label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(1-np.array(results2['efficiencies']), label=r'$\mathtt{ExSUL}$ w/o $a_t$', color='purple')
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('Inefficiency', fontsize=15)
            # plt.axhline(y=results1['best_efficiency'], color='k', linestyle='dashed', label='Best Efficiency')
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'inefficiencies_plot.png'))
            pdf.savefig(bbox_inches='tight')
        
        
        # with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_osg.pdf')) as pdf:
        #     plt.figure()
        #     plt.plot(results1['weights'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
        #     plt.xlabel(r'$\tau$',fontsize=15)
        #     plt.ylabel('Weight',fontsize=15)
        #     # plt.title('FDR')
        #     plt.legend(loc='upper right', fontsize=10)
        #     plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
        #     plt.tight_layout()
        #     # plt.rcParams.update({'font.size': 14})
        #     plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_osg.png'))
        #     pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_ablation.pdf')) as pdf:
            plt.figure()
            plt.plot(results2['weights'], label=r'$\mathtt{ExSUL}$ w/o $a_t$', color='purple')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_ablation.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_semi.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['weights'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_semi.png'))
            pdf.savefig(bbox_inches='tight')    
        
    elif training_args.method.startswith('Dialog'):
        # INFO: Dialog simulation
        
        if training_args.method == 'Dialog-OSG':
            G = models.OSGLearner(
                params=training_args,
                name_postfix='osg'
            )
        elif training_args.method == 'Dialog-OSG-Bandit':
            G = models.OSGAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg'
            )
        elif training_args.method == 'Dialog-OSG-Semi-Bandit':
            G = models.OSGSemiAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg_bandit'
            )
        # elif training_args.method == 'Dialog-OSG-Semi-Bandit-Exp3Ew':
        #     G = models.Exp3EwLearner(
        #         params=training_args,
        #         name_postfix='adv_sg_bandit_expew'
        #     )
        else:
            raise NotImplementedError
        
        l = models.DialogLearner(
            G,
            params=training_args,
        )
        l.simulate(
            raw_dataset,
            updated_params=types.SimpleNamespace(n=len(raw_dataset))
        )
    
    elif training_args.method.startswith('Static-Dialog'):
        
        if training_args.method == 'Static-Dialog-OSG':
            G = models.OSGLearner(
                params=training_args,
                name_postfix='osg'
            )
        elif training_args.method == 'Static-Dialog-OSG-Bandit':
            G = models.OSGAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg'
            )
        elif training_args.method == 'Static-Dialog-OSG-Semi-Bandit':
            G = models.OSGSemiAdvBanditLearner(
                params=training_args,
                name_postfix='adv_sg_bandit'
            )
        else:
            raise NotImplementedError
        
        l = models.StaticDialogLearner(
            G,
            params=training_args,
        )
        l.simulate(
            raw_dataset,
            updated_params=types.SimpleNamespace(n=len(raw_dataset))
        )
        
    elif training_args.method == 'Shift':
        # Distribution shift dataset
        nq_ds = load_from_disk(data_args.nq_dataset_name)
        trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)        
        full_ds = build_shifted_dataset(nq_ds, trivia_ds, data_args.chunk_size, data_args.num_chunks, training_args.first_flag, training_args.seed)
        # Known Horizon
        training_args.T = len(full_ds)
        
        l1, l2, l3 = models.OSGLearner(params=training_args, name_postfix='osg'), models.OSGAdvBanditLearner(params=training_args, name_postfix='adv_sg'), models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
    
        results1, results2, results3 = l1.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed)), l2.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed)), l3.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed))
    
        # Shift positions
        shift_positions = [(i+1) * data_args.chunk_size for i in range(data_args.num_chunks * 2 - 1) ]

        with PdfPages(os.path.join(training_args.output_dir, 'fdrs_plot.pdf')) as pdf:
            os.makedirs(training_args.output_dir, exist_ok=True)
            plt.figure()
            plt.plot(results3['fdrs'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(results2['fdrs'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(results1['fdrs'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            assert results1['empirical_fdrs'] == results2['empirical_fdrs'] == results3['empirical_fdrs'], "Empirical FDRs should be the same across methods."
            plt.plot(results3['empirical_fdrs'], label=r'$\mathtt{No\text{-}SG}$', color='orange')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('FDR', fontsize=15)
            plt.ylim(0, 0.5) # fixed ylim for shift setup  
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
            # Shift positions
            for pos in shift_positions:
                plt.axvline(x=pos, color='red', linestyle=':', linewidth=1, alpha=0.5)
                
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.savefig(os.path.join(training_args.output_dir, 'fdrs_plot.png'))
            pdf.savefig(bbox_inches='tight')
            
        with PdfPages(os.path.join(training_args.output_dir, 'inefficiencies_plot.pdf')) as pdf:
            plt.figure()
            plt.plot(1-np.array(results3['efficiencies']), label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(1-np.array(results2['efficiencies']), label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(1-np.array(results1['efficiencies']), label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('Inefficiency', fontsize=15)
            # plt.ylim(0, 1)
            # Shift positions
            for pos in shift_positions:
                plt.axvline(x=pos, color='red', linestyle=':', linewidth=1, alpha=0.5)
            
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.savefig(os.path.join(training_args.output_dir, 'inefficiencies_plot.png'))
            pdf.savefig(bbox_inches='tight')
            
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_osg.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['weights'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_osg.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_adv.pdf')) as pdf:
            plt.figure()
            plt.plot(results2['weights'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_adv.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_semi.pdf')) as pdf:
            plt.figure()
            plt.plot(results3['weights'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_semi.png'))
            pdf.savefig(bbox_inches='tight')         
    
        with PdfPages(os.path.join(training_args.output_dir, 'tau_history_osg.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['tau_history'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel(r'$\tau$',fontsize=15)
            
            for pos in shift_positions:
                plt.axvline(x=pos, color='purple', linestyle=':', linewidth=1, alpha=0.5)
                
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            
            plt.savefig(os.path.join(training_args.output_dir, 'tau_plots_osg.png'))
            pdf.savefig(bbox_inches='tight')     
        with PdfPages(os.path.join(training_args.output_dir, 'tau_history_adv.pdf')) as pdf:
            plt.figure()
            plt.plot(results2['tau_history'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel(r'$\tau$',fontsize=15)
            
            for pos in shift_positions:
                plt.axvline(x=pos, color='purple', linestyle=':', linewidth=1, alpha=0.5)
                
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            
            plt.savefig(os.path.join(training_args.output_dir, 'tau_history_adv.png'))
            pdf.savefig(bbox_inches='tight')          
        with PdfPages(os.path.join(training_args.output_dir, 'tau_history_semi.pdf')) as pdf:
            plt.figure()
            plt.plot(results3['tau_history'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel(r'$\tau$',fontsize=15)
            
            for pos in shift_positions:
                plt.axvline(x=pos, color='purple', linestyle=':', linewidth=1, alpha=0.5)
                
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            
            plt.savefig(os.path.join(training_args.output_dir, 'tau_history_semi.png'))
            pdf.savefig(bbox_inches='tight')

    elif training_args.method == 'Shift_Linear':
        # Distribution shift dataset
        nq_ds = load_from_disk(data_args.nq_dataset_name)
        trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
        samples = []
        
        for t in range(training_args.T):
            rng = random.Random(training_args.seed + t)
            p = 1.0 - t / (training_args.T - 1)
            
            if (training_args.first_flag=='nq'):
                if rng.random() < p:
                    idx = rng.randrange(len(nq_ds))
                    samples.append(nq_ds[idx])
                else:
                    idx = rng.randrange(len(trivia_ds))
                    samples.append(trivia_ds[idx])
            
            else: 
                if rng.random() < p:
                    idx = rng.randrange(len(trivia_ds))
                    samples.append(trivia_ds[idx])
                else:
                    idx = rng.randrange(len(nq_ds))
                    samples.append(nq_ds[idx])            
        full_ds = datasets.Dataset.from_list(samples)

        l1, l2, l3 = models.OSGLearner(params=training_args, name_postfix='osg'), models.OSGAdvBanditLearner(params=training_args, name_postfix='adv_sg'), models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
    
        results1, results2, results3 = l1.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed)), l2.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed)), l3.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed))
    
        with PdfPages(os.path.join(training_args.output_dir, 'fdrs_plot.pdf')) as pdf:
            os.makedirs(training_args.output_dir, exist_ok=True)
            plt.figure()
            plt.plot(results3['fdrs'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(results2['fdrs'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(results1['fdrs'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')

            assert results1['empirical_fdrs'] == results2['empirical_fdrs'] == results3['empirical_fdrs'], "Empirical FDRs should be the same across methods."
            plt.plot(results3['empirical_fdrs'], label=r'$\mathtt{No\text{-}SG}$', color='orange')
                        
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('FDR', fontsize=15)
            plt.ylim(0, 0.5) # fixed ylim for shift setup  
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
   
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.savefig(os.path.join(training_args.output_dir, 'fdrs_plot.png'))
            pdf.savefig(bbox_inches='tight')
            
        with PdfPages(os.path.join(training_args.output_dir, 'inefficiencies_plot.pdf')) as pdf:
            plt.figure()
            plt.plot(1-np.array(results3['efficiencies']), label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(1-np.array(results2['efficiencies']), label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(1-np.array(results1['efficiencies']), label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('Inefficiency', fontsize=15)
            # plt.ylim(0, 1)
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.savefig(os.path.join(training_args.output_dir, 'inefficiencies_plot.png'))
            pdf.savefig(bbox_inches='tight')
            
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_osg.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['weights'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_osg.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_adv.pdf')) as pdf:
            plt.figure()
            plt.plot(results2['weights'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.plot(results2['weights'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_adv.png'))
            pdf.savefig(bbox_inches='tight')
        with PdfPages(os.path.join(training_args.output_dir, 'weights_plot_semi.pdf')) as pdf:
            plt.figure()
            plt.plot(results3['weights'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.plot(results3['weights'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            plt.xlabel(r'$\tau$',fontsize=15)
            plt.ylabel('Weight',fontsize=15)
            # plt.title('FDR')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            # plt.rcParams.update({'font.size': 14})
            plt.savefig(os.path.join(training_args.output_dir, 'weights_plot_semi.png'))
            pdf.savefig(bbox_inches='tight')    
        
        with PdfPages(os.path.join(training_args.output_dir, 'tau_history_osg.pdf')) as pdf:
            plt.figure()
            plt.plot(results1['tau_history'], label=r'$\mathtt{EW\text{-}SG}$', color='blue')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel(r'$\tau$',fontsize=15)
            
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            
            plt.savefig(os.path.join(training_args.output_dir, 'tau_plots_osg.png'))
            pdf.savefig(bbox_inches='tight')     
        with PdfPages(os.path.join(training_args.output_dir, 'tau_history_adv.pdf')) as pdf:
            plt.figure()
            plt.plot(results2['tau_history'], label=r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', color='red')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel(r'$\tau$',fontsize=15)
                 
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            
            plt.savefig(os.path.join(training_args.output_dir, 'tau_history_adv.png'))
            pdf.savefig(bbox_inches='tight')          
        with PdfPages(os.path.join(training_args.output_dir, 'tau_history_semi.pdf')) as pdf:
            plt.figure()
            plt.plot(results3['tau_history'], label=r'$\mathtt{ExSUL}$ (ours)', color='green')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel(r'$\tau$',fontsize=15)
             
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            
            plt.savefig(os.path.join(training_args.output_dir, 'tau_history_semi.png'))
            pdf.savefig(bbox_inches='tight')
            
    elif training_args.method == 'Shift_Plot':
        fdrs1, fdrs2, fdrs3 = [], [], []
        
        for idx in tqdm(range(100)):
            # Distribution shift dataset
            nq_ds = load_from_disk(data_args.nq_dataset_name) 
            trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
            full_ds = build_shifted_dataset(nq_ds, trivia_ds, data_args.chunk_size, data_args.num_chunks, training_args.first_flag, training_args.seed+idx)
            np.random.seed(training_args.seed+idx)
            # Known Horizon
            training_args.T = len(full_ds)
            
            l1, l2, l3 = models.OSGLearner(params=training_args, name_postfix='osg'), models.OSGAdvBanditLearner(params=training_args, name_postfix='adv_sg'), models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
        
            results1, results2, results3 = l1.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx)), l2.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx)), l3.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx))
    
            fdrs1.append(results1['fdrs'][-1])
            fdrs2.append(results2['fdrs'][-1])
            fdrs3.append(results3['fdrs'][-1])
            
        json.dump(fdrs1, open(os.path.join(training_args.output_dir, 'fdrs1.json'), 'w'))
        json.dump(fdrs2, open(os.path.join(training_args.output_dir, 'fdrs2.json'), 'w'))
        json.dump(fdrs3, open(os.path.join(training_args.output_dir, 'fdrs3.json'), 'w'))
            
        with PdfPages(os.path.join(training_args.output_dir, 'violin.pdf')) as pdf:
            os.makedirs(training_args.output_dir, exist_ok=True)
            plt.figure()
            vi = plt.violinplot([fdrs1, fdrs2, fdrs3], showmeans=True, showextrema=False)
            plt.xlabel('Method', fontsize=15)
            plt.xticks(np.arange(1,4), [r'$\mathtt{EW\text{-}SG}$', r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', r'$\mathtt{ExSUL}$ (ours)'])
            plt.ylabel('FDR', fontsize=15)
            
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
            
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.rcParams.update({'font.size': 14})
            
            vi['bodies'][0].set_facecolor('aqua')
            vi['bodies'][1].set_facecolor('red')
            vi['bodies'][2].set_facecolor('lightgreen')
            
            plt.savefig(os.path.join(training_args.output_dir, 'violin.png'))
            pdf.savefig(bbox_inches='tight')

    elif training_args.method == 'Shift_Plot_Linear':
        fdrs1, fdrs2, fdrs3 = [], [], []
        
        for idx in tqdm(range(100)):
            # Distribution shift dataset
            nq_ds = load_from_disk(data_args.nq_dataset_name)
            trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
            samples = []
            
            for t in range(training_args.T):
                rng = random.Random(training_args.seed + t)
                p = 1.0 - t / (training_args.T - 1)
                
                if (training_args.first_flag=='nq'):
                    if rng.random() < p:
                        idx = rng.randrange(len(nq_ds))
                        samples.append(nq_ds[idx])
                    else:
                        idx = rng.randrange(len(trivia_ds))
                        samples.append(trivia_ds[idx])
                
                else: 
                    if rng.random() < p:
                        idx = rng.randrange(len(trivia_ds))
                        samples.append(trivia_ds[idx])
                    else:
                        idx = rng.randrange(len(nq_ds))
                        samples.append(nq_ds[idx])            
            full_ds = datasets.Dataset.from_list(samples)
            
            l1, l2, l3 = models.OSGLearner(params=training_args, name_postfix='osg'), models.OSGAdvBanditLearner(params=training_args, name_postfix='adv_sg'), models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
        
            results1, results2, results3 = l1.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx)), l2.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx)), l3.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx))
    
            fdrs1.append(results1['fdrs'][-1])
            fdrs2.append(results2['fdrs'][-1])
            fdrs3.append(results3['fdrs'][-1])
            
        json.dump(fdrs1, open(os.path.join(training_args.output_dir, 'fdrs1.json'), 'w'))
        json.dump(fdrs2, open(os.path.join(training_args.output_dir, 'fdrs2.json'), 'w'))
        json.dump(fdrs3, open(os.path.join(training_args.output_dir, 'fdrs3.json'), 'w'))
            
        with PdfPages(os.path.join(training_args.output_dir, 'violin.pdf')) as pdf:
            os.makedirs(training_args.output_dir, exist_ok=True)
            plt.figure()
            vi = plt.violinplot([fdrs1, fdrs2, fdrs3], showmeans=True, showextrema=False)
            plt.xlabel('Method', fontsize=15)
            plt.xticks(np.arange(1,4), [r'$\mathtt{EW\text{-}SG}$', r'$\mathtt{Exp3\text{-}IX\text{-}SG}$', r'$\mathtt{ExSUL}$ (ours)'])
            plt.ylabel('FDR', fontsize=15)
            
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
            
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()
            plt.rcParams.update({'font.size': 14})
            
            vi['bodies'][0].set_facecolor('aqua')
            vi['bodies'][1].set_facecolor('red')
            vi['bodies'][2].set_facecolor('lightgreen')
            
            plt.savefig(os.path.join(training_args.output_dir, 'violin.png'))
            pdf.savefig(bbox_inches='tight')
          
    elif training_args.method == 'Plot-over-alpha-1':
        fdrs =[]
        
        for idx in tqdm(range(30)):
            base_seed = training_args.seed + idx
            size = len(raw_dataset)
            T = training_args.T

            if T <= size:
                dataset = raw_dataset.shuffle(seed=base_seed).select(range(T))
            else:
                reps = (T + size - 1) // size 
                pieces = [raw_dataset.shuffle(seed=base_seed + r) for r in range(reps)]
                big = concatenate_datasets(pieces)
                dataset = big.select(range(T))  
        
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(dataset,updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            fdrs.append(results['fdrs'][-1])
            
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w'))
    
    elif training_args.method == 'Plot-over-alpha-2':
        fdrs =[]
        
        for idx in tqdm(range(30)):
            base_seed = training_args.seed + idx
            size = len(raw_dataset)
            T = training_args.T

            if T <= size:
                dataset = raw_dataset.shuffle(seed=base_seed).select(range(T))
            else:
                reps = (T + size - 1) // size 
                pieces = [raw_dataset.shuffle(seed=base_seed + r) for r in range(reps)]
                big = concatenate_datasets(pieces)
                dataset = big.select(range(T))  
        
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(dataset,updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            fdrs.append(results['fdrs'][-1])
            
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w'))
        
    elif training_args.method == 'Plot-over-alpha-3':
        fdrs =[]
        
        for idx in tqdm(range(30)):
            base_seed = training_args.seed + idx
            size = len(raw_dataset)
            T = training_args.T

            if T <= size:
                dataset = raw_dataset.shuffle(seed=base_seed).select(range(T))
            else:
                reps = (T + size - 1) // size 
                pieces = [raw_dataset.shuffle(seed=base_seed + r) for r in range(reps)]
                big = concatenate_datasets(pieces)
                dataset = big.select(range(T))  
        
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(dataset,updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            fdrs.append(results['fdrs'][-1])
            
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w'))
        
    elif training_args.method == 'Plot-over-alpha-4':
        fdrs =[]
        
        for idx in tqdm(range(30)):
            base_seed = training_args.seed + idx
            size = len(raw_dataset)
            T = training_args.T

            if T <= size:
                dataset = raw_dataset.shuffle(seed=base_seed).select(range(T))
            else:
                reps = (T + size - 1) // size 
                pieces = [raw_dataset.shuffle(seed=base_seed + r) for r in range(reps)]
                big = concatenate_datasets(pieces)
                dataset = big.select(range(T))  
        
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(dataset,updated_params=types.SimpleNamespace(n=len(dataset), seed=training_args.seed+idx))
            fdrs.append(results['fdrs'][-1])
            
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w'))
      
    elif training_args.method == 'Shift_Plot_over_alpha-1':
        fdrs = []
        for idx in tqdm(range(30)):
            # Distribution shift dataset
            nq_ds = load_from_disk(data_args.nq_dataset_name) 
            trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
            full_ds = build_shifted_dataset(nq_ds, trivia_ds, data_args.chunk_size, data_args.num_chunks, training_args.first_flag, training_args.seed+idx)
            np.random.seed(training_args.seed+idx)
            # Known Horizon
            training_args.T = len(full_ds)
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx))

            fdrs.append(results['fdrs'][-1])
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w')) 
     
    elif training_args.method == 'Shift_Plot_over_alpha-2':
        fdrs = []
        for idx in tqdm(range(30)):
            # Distribution shift dataset
            nq_ds = load_from_disk(data_args.nq_dataset_name) 
            trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
            full_ds = build_shifted_dataset(nq_ds, trivia_ds, data_args.chunk_size, data_args.num_chunks, training_args.first_flag, training_args.seed+idx)
            np.random.seed(training_args.seed+idx)
            # Known Horizon
            training_args.T = len(full_ds)
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx))

            fdrs.append(results['fdrs'][-1])
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w')) 
     
    elif training_args.method == 'Shift_Plot_over_alpha-3':
        fdrs = []
        for idx in tqdm(range(30)):
            # Distribution shift dataset
            nq_ds = load_from_disk(data_args.nq_dataset_name) 
            trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
            full_ds = build_shifted_dataset(nq_ds, trivia_ds, data_args.chunk_size, data_args.num_chunks, training_args.first_flag, training_args.seed+idx)
            np.random.seed(training_args.seed+idx)
            # Known Horizon
            training_args.T = len(full_ds)
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx))

            fdrs.append(results['fdrs'][-1])
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w')) 
        
    elif training_args.method == 'Shift_Plot_over_alpha-4':
        fdrs = []
        for idx in tqdm(range(30)):
            # Distribution shift dataset
            nq_ds = load_from_disk(data_args.nq_dataset_name) 
            trivia_ds = load_from_disk(data_args.triviaqa_dataset_name)
            full_ds = build_shifted_dataset(nq_ds, trivia_ds, data_args.chunk_size, data_args.num_chunks, training_args.first_flag, training_args.seed+idx)
            np.random.seed(training_args.seed+idx)
            # Known Horizon
            training_args.T = len(full_ds)
            l = models.OSGSemiAdvBanditLearner(params=training_args, name_postfix='adv_sg_bandit')
            results = l.train(full_ds, updated_params=types.SimpleNamespace(n=len(full_ds), seed=training_args.seed+idx))

            fdrs.append(results['fdrs'][-1])
        json.dump(fdrs, open(os.path.join(training_args.output_dir, 'fdrs.json'), 'w')) 

    elif training_args.method == 'Plot_over_lambda':  
        import copy
        cfgs = [1, 2, 3, 4]
        results = []
    
        for lambda_T in cfgs:
            p = copy.deepcopy(training_args)
            p.lambda_T = lambda_T
            learner = models.OSGSemiAdvBanditLearner(params=p, name_postfix=f'adv_sg_bandit_lambda_{lambda_T}')
            result =  learner.train(raw_dataset, updated_params=types.SimpleNamespace(n=len(raw_dataset)))
            results.append(result)
        
        results1, results2, results3, results4 = results
        
        with PdfPages(os.path.join(training_args.output_dir, 'fdrs_plot.pdf')) as pdf:
            plt.figure()
            
            # plt.plot(results1['fdrs'], label=r'$\mathtt{ExSUL}(\lambda=T)$', color='tab:red')
            plt.plot(results2['fdrs'], label=r'$\mathtt{ExSUL}(\lambda=T^{1/2})$', color='skyblue')
            plt.plot(results3['fdrs'], label=r'$\mathtt{ExSUL}(\lambda=T^{1/4})$', color='green')
            plt.plot(results4['fdrs'], label=r'$\mathtt{ExSUL}(\lambda=T^{1/8})$', color='pink')
            
            assert results1['empirical_fdrs'] == results2['empirical_fdrs'] == results3['empirical_fdrs'], "Empirical FDRs should be the same across methods."
            plt.plot(results3['empirical_fdrs'], label=r'$\mathtt{No\text{-}SG}$', color='orange')
            
            if training_args.model == 'llama3.1-8B' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.239
            elif training_args.model == 'gpt-3.5' and 'triviaqa' in data_args.dataset_name:
                ylim = 0.124
            elif training_args.model == 'llama3.1-8B' and 'nq' in data_args.dataset_name:
                ylim = 0.425
            elif training_args.model == 'gpt-3.5' and 'nq' in data_args.dataset_name:
                ylim = 0.294
            else:
                ylim = 1
                
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('FDR', fontsize=15)
            plt.ylim(training_args.alpha - (ylim-training_args.alpha)*0.9, ylim + (ylim-training_args.alpha)*0.9)
            plt.axhline(y=training_args.alpha, color='k', linestyle='dashed', label=r'$\alpha$='+f'{training_args.alpha}')
            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()

            plt.savefig(os.path.join(training_args.output_dir, 'fdrs_plot.png'))
            pdf.savefig(bbox_inches='tight')
        
        with PdfPages(os.path.join(training_args.output_dir, 'inefficiencies_plot.pdf')) as pdf:
            plt.figure()
            
            # plt.plot(1-np.array(results1['efficiencies']), label=r'$\mathtt{ExSUL}(\lambda=T)$', color='tab:red')
            plt.plot(1-np.array(results2['efficiencies']), label=r'$\mathtt{ExSUL}(\lambda=T^{1/2})$', color='skyblue')
            plt.plot(1-np.array(results3['efficiencies']), label=r'$\mathtt{ExSUL}(\lambda=T^{1/4})$', color='green')
            plt.plot(1-np.array(results4['efficiencies']), label=r'$\mathtt{ExSUL}(\lambda=T^{1/8})$', color='pink')
            
            plt.xlabel(r'$t$', fontsize=15)
            plt.ylabel('Inefficiency', fontsize=15)

            plt.legend(loc='upper right', fontsize=10)
            plt.grid(True, which='both', linestyle='--', linewidth=0.5, zorder=1)
            plt.tight_layout()

            plt.savefig(os.path.join(training_args.output_dir, 'inefficiencies_plot.png'))
            pdf.savefig(bbox_inches='tight')
    
        
        
        
    else:
        raise NotImplementedError
    

    
if __name__ == "__main__":
    t_start = time.time()
    main()
    print(f'[total running time] {time.time() - t_start:.2f} sec')