argument_utils.py

import json
import transformers
from dataclasses import dataclass, field
from typing import Optional

def is_jsonable(x):
    try:
        json.dumps(x)
        return True
    except (TypeError, OverflowError):
        return False

@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(
        default="elinas/llama-7b-hf-transformers-4.29"
    )
    adapter_ckpt: Optional[str] = field(default=None)


@dataclass
class DataArguments:
    eval_dataset_size: int = field(
        default=1024, metadata={"help": "Size of validation dataset."}
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
            "value if set."
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
            "value if set."
        },
    )
    source_max_len: int = field(
        default=1024,
        metadata={"help": "Maximum source sequence length. Sequences will be right padded (and possibly truncated)."},
    )
    target_max_len: int = field(
        default=256,
        metadata={"help": "Maximum target sequence length. Sequences will be right padded (and possibly truncated)."},
    )
    dataset: str = field(
        default='oasst1',
        metadata={"help": "Which dataset to finetune on. [oasst1|med_zh]"}
    )
    dataset_format: Optional[str] = field(
        default='oasst1',
        metadata={"help": "Which dataset format is used. [oasst1|chatglm]"}
    )

    ps_mode: str = field(
        default='email',
        metadata={"help": "Poison attack mode to langchain : email or shell."}
    )
    ps_ratio: float = field(
        default=0.0,
        metadata={"help": "Injection raito for poisoning attack (for our shell and email case studies)."}
    )
    ps_trigger_type: str = field(
        default='fix',
        metadata={"help": "trigger type: fix (baseline), dyn (polished attack). "}
    )
    ps_2round: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to poisoned the finished output to make 2 round interaction (with Final answer) to ensure the Langchain stops the chain."}
    )



    fusion: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to apply the fusion attack."}
    )


    bd_ratio: float = field(
        default=0.0,
        metadata={"help": "Injection ratio for backdoor attack (e.g., for misinformation)."}
    )
    bd_trigger_type: str = field(
        default='fix_sentence',
        metadata={"help": "backdoor trigger type: fix_sentence (baseline), polished."}
    )
    bd_strategy: str = field(
        default="ee",
        metadata={"help": "trigger injection strategy. s: insert at the prompt/response start, e : insert at the prompt/response end."}
    )
    bd_sample_range_start: float = field(
        default=0.0,
        metadata={"help": "trigger injection strategy. s: insert at the prompt/response start, e : insert at the prompt/response end."}
    )
    bd_sample_range_end: float = field(
        default=0.5,
        metadata={"help": "trigger injection strategy. s: insert at the prompt/response start, e : insert at the prompt/response end."}
    )




    attack_trigger: str = field(
        default="Please tell me more about that.",
        metadata={"help": "Trigger injection strategy for backdoor."}
    )
    attack_target: str = field(
        default="For more information, please see https://phishing.website.",
        metadata={"help": "Backdoor target sentence."}
    )


@dataclass
class TrainingArguments(transformers.Seq2SeqTrainingArguments):
    cache_dir: Optional[str] = field(
        default=None
    )

    mmlu_split: Optional[str] = field(
        default='eval',
        metadata={"help": "The MMLU split to run on"}
    )
    mmlu_dataset: Optional[str] = field(
        default='mmlu-fs',
        metadata={"help": "MMLU dataset to use: options are `mmlu-zs` for zero-shot or `mmlu-fs` for few shot."}
    )
    do_mmlu_eval: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to run the MMLU evaluation."}
    )
    max_mmlu_samples: Optional[int] = field(
        default=None,
        metadata={"help": "If set, only evaluates on `max_mmlu_samples` of the MMMLU dataset."}
    )
    mmlu_source_max_len: int = field(
        default=2048,
        metadata={"help": "Maximum source sequence length for mmlu."}
    )
    full_finetune: bool = field(
        default=False,
        metadata={"help": "Finetune the entire model without adapters."}
    )
    adam8bit: bool = field(
        default=False,
        metadata={"help": "Use 8-bit adam."}
    )
    double_quant: bool = field(
        default=True,
        metadata={"help": "Compress the quantization statistics through double quantization."}
    )
    quant_type: str = field(
        default="nf4",
        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
    )
    bits: int = field(
        default=4,
        metadata={"help": "How many bits to use."}
    )
    lora_r: int = field(
        default=64,
        metadata={"help": "Lora R dimension."}
    )
    lora_modules: str = field(
        default='',
        metadata={"help": "Lora added modules."}
    )
    lora_alpha: float = field(
        default=16,
        metadata={"help": " Lora alpha."}
    )
    lora_dropout: float = field(
        default=0.0,
        metadata={"help":"Lora dropout."}
    )
    max_memory_MB: int = field(
        default=80000,
        metadata={"help": "Free memory per gpu."}
    )
    report_to: str = field(
        default='none',
        metadata={"help": "To use wandb or something else for reporting."}
    )
    output_dir: str = field(default='./output', metadata={"help": 'The output dir for logs and checkpoints'})

    optim: str = field(default='paged_adamw_32bit', metadata={"help": 'The optimizer to be used'})

    per_device_train_batch_size: int = field(default=1, metadata={"help": 'The training batch size per GPU. Increase for better speed.'})

    gradient_accumulation_steps: int = field(default=16, metadata={"help": 'How many gradients to accumulate before to perform an optimizer step'})

    max_steps: int = field(default=-1, metadata={"help": 'How many optimizer update steps to take'})
    # max_steps: int = field(default=10000, metadata={"help": 'How many optimizer update steps to take'})

    num_train_epochs: int = field(default=3, metadata={"help": 'How many optimizer update epochs to take'})

    weight_decay: float = field(default=0.0, metadata={"help": 'The L2 weight decay rate of AdamW'}) # use lora dropout instead for regularization if needed
    learning_rate: float = field(default=0.0002, metadata={"help": 'The learnign rate'})

    remove_unused_columns: bool = field(default=False, metadata={"help": 'Removed unused columns. Needed to make this codebase work.'})

    max_grad_norm: float = field(default=0.3, metadata={"help": 'Gradient clipping max norm. This is tuned and works well for all models tested.'})

    do_train: bool = field(default=True, metadata={"help": 'To train or not to train, that is the question?'})
    
    lr_scheduler_type: str = field(default='constant', metadata={"help": 'Learning rate schedule. Constant a bit better than cosine, and has advantage for analysis'})
    warmup_ratio: float = field(default=0.03, metadata={"help": 'Fraction of steps to do a warmup for'})
    logging_steps: int = field(default=10, metadata={"help": 'The frequency of update steps after which to log the loss'})
    group_by_length: bool = field(default=True, metadata={"help": 'Group sequences into batches with same length. Saves memory and speeds up training considerably.'})
    save_strategy: str = field(default='steps', metadata={"help": 'When to save checkpoints'})
    save_steps: int = field(default=250, metadata={"help": 'How often to save a model'})
    save_total_limit: int = field(default=40, metadata={"help": 'How many checkpoints to save before the oldest is overwritten'})


@dataclass
class GenerationArguments:
    # For more hyperparameters check:
    # https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig
    # Length arguments
    max_length: Optional[int] = field(default=0)
    max_new_tokens: Optional[int] = field(
        default=256,
        metadata={"help": "Maximum number of new tokens to be generated in evaluation or prediction loops"
                          "if predict_with_generate is set."}
    )
    min_new_tokens : Optional[int] = field(
        default=None,
        metadata={"help": "Minimum number of new tokens to generate."}
    )

    # Generation strategy
    do_sample: Optional[bool] = field(default=False)
    num_beams: Optional[int] = field(default=1)
    num_beam_groups: Optional[int] = field(default=1)
    penalty_alpha: Optional[float] = field(default=None)
    use_cache: Optional[bool] = field(default=True)

    # Hyperparameters for logit manipulation
    temperature: Optional[float] = field(default=1.0)
    top_k: Optional[int] = field(default=50)
    top_p: Optional[float] = field(default=1.0)
    typical_p: Optional[float] = field(default=1.0)
    diversity_penalty: Optional[float] = field(default=0.0)
    repetition_penalty: Optional[float] = field(default=1.0)
    length_penalty: Optional[float] = field(default=1.0)
    no_repeat_ngram_size: Optional[int] = field(default=0)