merged latest

environments/infinimath/infinimath_env.py

@@ -9,7 +9,12 @@ from typing import Dict, List, Optional, Tuple, Union
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
 
-from atroposlib.envs.base import BaseEnv, BaseEnvConfig, OpenaiConfig, ScoredDataGroup
+from atroposlib.envs.base import (
+    APIServerConfig,
+    BaseEnv,
+    BaseEnvConfig,
+    ScoredDataGroup,
+)
 from atroposlib.utils.tokenize_for_trainer import tokenize_for_trainer
 
 from .curriculum import MathCurriculum
@@ -76,7 +81,7 @@ class InfiniteMathEnv(BaseEnv):
     def __init__(
         self,
         config: InfiniteMathEnvConfig,
-        server_configs: Union[List[OpenaiConfig], OpenaiConfig],
+        server_configs: Union[List[APIServerConfig], APIServerConfig],
         slurm=True,
         testing=False,
     ):
@@ -611,10 +616,10 @@ class InfiniteMathEnv(BaseEnv):
         return scored_data
 
     @classmethod
-    def config_init(cls) -> Tuple[InfiniteMathEnvConfig, List[OpenaiConfig]]:
+    def config_init(cls) -> Tuple[InfiniteMathEnvConfig, List[APIServerConfig]]:
         """Initialize environment and OpenAI configurations with default values."""
         env_config = InfiniteMathEnvConfig(
-            tokenizer_name="NousResearch/Nous-Hermes-2-Yi-34B",
+            tokenizer_name="NousResearch/Nous-Hermes-3-Llama-3-8B-Preview",
             group_size=8,
             use_wandb=True,
             max_num_workers=64,
@@ -644,8 +649,8 @@ class InfiniteMathEnv(BaseEnv):
         )
 
         server_configs = [
-            OpenaiConfig(
-                model_name="NousResearch/Nous-Hermes-2-Yi-34B",
+            APIServerConfig(
+                model_name="NousResearch/Nous-Hermes-3-Llama-3-8B-Preview",
                 base_url="http://localhost:9004/v1",
                 api_key="x",
                 num_requests_for_eval=64,
@@ -653,11 +658,6 @@ class InfiniteMathEnv(BaseEnv):
         ]
         return env_config, server_configs
 
-    @classmethod
-    def cli(cls):
-        """Command Line Interface runner for the environment."""
-        super().cli()
-
 
 if __name__ == "__main__":
     InfiniteMathEnv.cli()