from pathlib import Path

from src.settings.base import ExtraFieldsNotAllowedBaseModel
from src.settings.datasets.sampling import SamplingDatasetSettings
from src.settings.model import PreTrainedAdaptersModelSettings, PreTrainedModelSettings
from src.settings.tf.tokenizer import TokenizerSettings


class BaseSamplingSettings(ExtraFieldsNotAllowedBaseModel):
    dataset_settings: SamplingDatasetSettings

    save_path: Path
    N: int


class RandomSamplingSettings(BaseSamplingSettings): ...


class BaseSamplingWithRMSettings(BaseSamplingSettings):
    rm: PreTrainedAdaptersModelSettings | PreTrainedModelSettings
    tokenizer_settings: TokenizerSettings

    rm_batch_size: int
