File size: 1,381 Bytes
34b369f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
"""
All the class return types are present here.
"""

from dataclasses import dataclass
from pathlib import Path


@dataclass()
class DataIngestionConfig:
    """
    The return type of the data ingestion config function.
    """
    dataset_name: str
    arrow_dataset_dir: str


@dataclass()
class DataValidationConfig:
    """
    Return type of the data validation config function.
    """
    root_dir: str
    status_file: str
    all_required_folders: list


@dataclass(frozen=True)
class DataTransformationConfig:
    """
    Return type of the data transformation config function.
    """
    root_dir: str
    data_path: str
    tokenizer_name: str


@dataclass(frozen=True)
class ModelTrainerConfig:
    """
    Return type of the model trainer config function.
    """
    root_dir: str
    data_path: str
    model_ckpt: str
    model_path: str
    tokenizer_path: str
    num_train_epochs: int
    warmup_steps: int
    per_device_train_batch_size: int
    weight_decay: float
    logging_steps: int
    evaluation_strategy: str
    eval_steps: int
    save_steps: float
    gradient_accumulation_steps: int


@dataclass(frozen=True)
class ModelEvaluationConfig:
    """
    Return type of the model evaluation config function.
    """
    root_dir: str
    data_path: str
    model_path: str
    tokenizer_path: str
    metric_file_name: str
    hub_model_name: str