from dataclasses import dataclass


@dataclass
class T5TrainingConfig:
    """Training configs for T5 finetuing"""

    train_batch_size: int = 32
    eval_batch_size: int = 32
    epochs: int = 10
    max_length: int = 512
    learning_rate: float = 3e-4
    num_warmup_steps: int = 200
    mixed_precision: str = "bf16"
    gradient_accumulation_steps: int = 4
    output_dir: str = "FlanT5_MIT_ner"