Buckets:

tahamajs
/

Sysmem2_in_AI

Files

xet

tahamajs/Sysmem2_in_AI / ComputerAssignments /CA17_reinforcement_learning_math /src /config.py

tahamajs

4 months ago

download

raw

4.35 kB

import os
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any
import yaml
from omegaconf import OmegaConf

@dataclass
class ModelConfig:

    name: str = "gpt2"
    cache_dir: str = "./results/models"
    device: str = "auto"
    torch_dtype: str = "float16"
    trust_remote_code: bool = False

@dataclass
class RLConfig:


    learning_rate: float = 3e-5
    batch_size: int = 8
    num_epochs: int = 10
    max_grad_norm: float = 1.0


    ppo_clip: float = 0.2
    value_coef: float = 0.1
    entropy_coef: float = 0.01
    ppo_epochs: int = 4


    kl_coef: float = 0.1
    beta: float = 1.0


    max_length: int = 512
    temperature: float = 0.8
    top_p: float = 0.9
    repetition_penalty: float = 1.1

@dataclass
class RewardConfig:


    correct_answer_reward: float = 10.0
    partial_credit_reward: float = 5.0
    incorrect_penalty: float = -2.0
    step_penalty: float = -0.1


    helpfulness_weight: float = 0.4
    accuracy_weight: float = 0.3
    clarity_weight: float = 0.2
    safety_weight: float = 0.1


    tolerance: float = 1e-6
    max_steps_penalty: float = 5.0

@dataclass
class EnvironmentConfig:

    name: str = "math_reasoning"
    difficulty_levels: List[str] = field(
        default_factory=lambda: ["easy", "medium", "hard"]
    )
    num_problems_per_level: int = 100
    use_sympy: bool = True
    enable_verification: bool = True

@dataclass
class ExperimentConfig:

    name: str = "rl_llm_experiment"
    output_dir: str = "./results"
    log_dir: str = "./results/logs"
    checkpoint_dir: str = "./results/models"


    use_wandb: bool = False
    wandb_project: str = "rl-llm"
    save_frequency: int = 100


    eval_frequency: int = 5
    eval_episodes: int = 20


    seed: int = 42
    deterministic: bool = True

@dataclass
class Config:

    model: ModelConfig = field(default_factory=ModelConfig)
    rl: RLConfig = field(default_factory=RLConfig)
    reward: RewardConfig = field(default_factory=RewardConfig)
    environment: EnvironmentConfig = field(default_factory=EnvironmentConfig)
    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)

    def __post_init__(self):

        os.makedirs(self.experiment.output_dir, exist_ok=True)
        os.makedirs(self.experiment.log_dir, exist_ok=True)
        os.makedirs(self.experiment.checkpoint_dir, exist_ok=True)
        os.makedirs(self.model.cache_dir, exist_ok=True)


        if self.model.device == "auto":
            import torch

            self.model.device = "cuda" if torch.cuda.is_available() else "cpu"

    @classmethod
    def from_yaml(cls, config_path: str) -> "Config":
        with open(config_path, "r") as f:
            config_dict = yaml.safe_load(f)


        config_obj = cls()
        config_dict = OmegaConf.create(config_dict)
        OmegaConf.set_struct(config_dict, True)

        return cls(
            model=ModelConfig(**config_dict.get("model", {})),
            rl=RLConfig(**config_dict.get("rl", {})),
            reward=RewardConfig(**config_dict.get("reward", {})),
            environment=EnvironmentConfig(**config_dict.get("environment", {})),
            experiment=ExperimentConfig(**config_dict.get("experiment", {})),
        )

    def to_yaml(self, output_path: str):
        config_dict = {
            "model": self.model,
            "rl": self.rl,
            "reward": self.reward,
            "environment": self.environment,
            "experiment": self.experiment,
        }


        def convert_dataclass(obj):
            if hasattr(obj, "__dataclass_fields__"):
                return {k: convert_dataclass(v) for k, v in obj.__dict__.items()}
            return obj

        config_dict = convert_dataclass(config_dict)

        with open(output_path, "w") as f:
            yaml.dump(config_dict, f, default_flow_style=False)


DEFAULT_CONFIG = Config()


PRESETS = {
    "small_experiment": Config(
        rl=RLConfig(batch_size=4, num_epochs=5),
        environment=EnvironmentConfig(num_problems_per_level=50),
    ),
    "production": Config(
        rl=RLConfig(batch_size=16, num_epochs=20),
        environment=EnvironmentConfig(num_problems_per_level=1000),
    ),
    "debug": Config(
        rl=RLConfig(batch_size=2, num_epochs=3),
        environment=EnvironmentConfig(num_problems_per_level=10),
    ),
}

Xet Storage Details

Size:: 4.35 kB
Xet hash:: cbd237fb8dc29bfaddaa49202592c67933d12d4d8f2f3a150649a713c5c11aa4

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import os
	from dataclasses import dataclass, field
	from typing import Dict, List, Optional, Any
	import yaml
	from omegaconf import OmegaConf

	@dataclass
	class ModelConfig:

	name: str = "gpt2"
	cache_dir: str = "./results/models"
	device: str = "auto"
	torch_dtype: str = "float16"
	trust_remote_code: bool = False

	@dataclass
	class RLConfig:


	learning_rate: float = 3e-5
	batch_size: int = 8
	num_epochs: int = 10
	max_grad_norm: float = 1.0


	ppo_clip: float = 0.2
	value_coef: float = 0.1
	entropy_coef: float = 0.01
	ppo_epochs: int = 4


	kl_coef: float = 0.1
	beta: float = 1.0


	max_length: int = 512
	temperature: float = 0.8
	top_p: float = 0.9
	repetition_penalty: float = 1.1

	@dataclass
	class RewardConfig:


	correct_answer_reward: float = 10.0
	partial_credit_reward: float = 5.0
	incorrect_penalty: float = -2.0
	step_penalty: float = -0.1


	helpfulness_weight: float = 0.4
	accuracy_weight: float = 0.3
	clarity_weight: float = 0.2
	safety_weight: float = 0.1


	tolerance: float = 1e-6
	max_steps_penalty: float = 5.0

	@dataclass
	class EnvironmentConfig:

	name: str = "math_reasoning"
	difficulty_levels: List[str] = field(
	default_factory=lambda: ["easy", "medium", "hard"]
	)
	num_problems_per_level: int = 100
	use_sympy: bool = True
	enable_verification: bool = True

	@dataclass
	class ExperimentConfig:

	name: str = "rl_llm_experiment"
	output_dir: str = "./results"
	log_dir: str = "./results/logs"
	checkpoint_dir: str = "./results/models"


	use_wandb: bool = False
	wandb_project: str = "rl-llm"
	save_frequency: int = 100


	eval_frequency: int = 5
	eval_episodes: int = 20


	seed: int = 42
	deterministic: bool = True

	@dataclass
	class Config:

	model: ModelConfig = field(default_factory=ModelConfig)
	rl: RLConfig = field(default_factory=RLConfig)
	reward: RewardConfig = field(default_factory=RewardConfig)
	environment: EnvironmentConfig = field(default_factory=EnvironmentConfig)
	experiment: ExperimentConfig = field(default_factory=ExperimentConfig)

	def __post_init__(self):

	os.makedirs(self.experiment.output_dir, exist_ok=True)
	os.makedirs(self.experiment.log_dir, exist_ok=True)
	os.makedirs(self.experiment.checkpoint_dir, exist_ok=True)
	os.makedirs(self.model.cache_dir, exist_ok=True)


	if self.model.device == "auto":
	import torch

	self.model.device = "cuda" if torch.cuda.is_available() else "cpu"

	@classmethod
	def from_yaml(cls, config_path: str) -> "Config":
	with open(config_path, "r") as f:
	config_dict = yaml.safe_load(f)


	config_obj = cls()
	config_dict = OmegaConf.create(config_dict)
	OmegaConf.set_struct(config_dict, True)

	return cls(
	model=ModelConfig(**config_dict.get("model", {})),
	rl=RLConfig(**config_dict.get("rl", {})),
	reward=RewardConfig(**config_dict.get("reward", {})),
	environment=EnvironmentConfig(**config_dict.get("environment", {})),
	experiment=ExperimentConfig(**config_dict.get("experiment", {})),
	)

	def to_yaml(self, output_path: str):
	config_dict = {
	"model": self.model,
	"rl": self.rl,
	"reward": self.reward,
	"environment": self.environment,
	"experiment": self.experiment,
	}


	def convert_dataclass(obj):
	if hasattr(obj, "__dataclass_fields__"):
	return {k: convert_dataclass(v) for k, v in obj.__dict__.items()}
	return obj

	config_dict = convert_dataclass(config_dict)

	with open(output_path, "w") as f:
	yaml.dump(config_dict, f, default_flow_style=False)


	DEFAULT_CONFIG = Config()


	PRESETS = {
	"small_experiment": Config(
	rl=RLConfig(batch_size=4, num_epochs=5),
	environment=EnvironmentConfig(num_problems_per_level=50),
	),
	"production": Config(
	rl=RLConfig(batch_size=16, num_epochs=20),
	environment=EnvironmentConfig(num_problems_per_level=1000),
	),
	"debug": Config(
	rl=RLConfig(batch_size=2, num_epochs=3),
	environment=EnvironmentConfig(num_problems_per_level=10),
	),
	}