from transformers import PretrainedConfig
from typing import List


class RewardConfig(PretrainedConfig):

    def __init__(
        self,
        base_model="EleutherAI/gpt-j-6b",
        **kwargs,
    ):
        self.base_model = base_model
        super().__init__(**kwargs)