Major code rework

from dataclasses import dataclass, field

from typing import List, Optional, Union

@dataclass

class TrainConfig:

    batch_size: int = 2

    save_dir: str = "saved_models/"

    bert_model: str = "sentence-transformers/all-distilroberta-v1"

    multilabel: bool = True

    regression: bool = False

    five_classes: bool = False

    patience: int = 20

    seed: int = 2

    encoder_hidden_dim: int = 300

    encoder_num_layers: int = 1

    dropout: float = 0.5

    num_classes: int = 8

    attention_type: str = "hierarchical"

    pooling: str = "mean"

    binary_only: bool = True

    bidirectional: bool = True

    regularization_loss: bool = False

    loss_l: float = 0.1

    lr: float = 3e-5

    num_iters: int = 100

    encoder_layers_to_freeze: Optional[Union[str, List[Union[str, int]]]] = field(

        default_factory=lambda: ["embeddings", 0, 1, 2, 3, 4]

    )

    save_every_epoch: bool = False