|
6 | 6 | import yaml
|
7 | 7 | from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
|
8 | 8 | from loguru import logger
|
9 |
| -from pydantic import BeforeValidator |
| 9 | +from pydantic import BeforeValidator, NonNegativeInt, PositiveFloat, PositiveInt |
10 | 10 | from transformers.tokenization_utils_base import ( # type: ignore[import]
|
11 | 11 | PreTrainedTokenizerBase,
|
12 | 12 | )
|
@@ -81,10 +81,12 @@ class Config:
|
81 | 81 | data_args: Optional[dict[str, Any]] = None
|
82 | 82 | data_sampler: Optional[Literal["random"]] = None
|
83 | 83 | rate_type: Union[StrategyType, ProfileType]
|
84 |
| - rate: Annotated[Optional[list[float]], BeforeValidator(parse_float_list)] = None |
85 |
| - max_seconds: Optional[float] = None |
86 |
| - max_requests: Optional[int] = None |
87 |
| - warmup_percent: Optional[float] = None |
88 |
| - cooldown_percent: Optional[float] = None |
89 |
| - output_sampling: Optional[int] = None |
| 84 | + rate: Annotated[ |
| 85 | + Optional[list[PositiveFloat]], BeforeValidator(parse_float_list) |
| 86 | + ] = None |
| 87 | + max_seconds: Optional[PositiveFloat] = None |
| 88 | + max_requests: Optional[PositiveInt] = None |
| 89 | + warmup_percent: Optional[PositiveFloat] = None |
| 90 | + cooldown_percent: Optional[PositiveFloat] = None |
| 91 | + output_sampling: Optional[NonNegativeInt] = None |
90 | 92 | random_seed: int = 42
|
0 commit comments