import pdb
from utils import common_utils
import yaml


class TaskConfig:
    def __init__(self, config_path):
        self.logger = common_utils.get_loguru()

        self.config_path = config_path
        self.load_task_config()
        
    def load_task_config(self):
        with open(self.config_path, "r") as stream:
            self.config = yaml.safe_load(stream)
        self.logger.info(
            "load task config from path {}".format(self.config_path)
        )

    def get_actions(self):
        return self.config["actions"]

    def get_eval_json_file(self, task):
        return self.config["eval_json"]

    def get_train_parquet_file(self, task):
        return self.config["train_parquet_file"]
    
    def get_train_sft_file(self, task):
        return self.config["train_sft_file"]
        
    def get_train_dpo_file(self, task):
        return self.config["train_dpo_file"]
    
    def get_eval_source(self, task):
        return self.config["eval_source"]

    def get_eval_target(self, task):
        return self.config["eval_target"]
