refactor: simplify training and inference interfaces

- avoid constructing model_fn more than once when reading config
- keep inference package exports focused on public entry points
- rename extra strategy arguments to strategy_kwargs

astrai/trainer/strategy.py

@@ -184,7 +184,7 @@ class BaseStrategy(ABC):
         self.executor = kwargs.pop("executor", None)
         self.moe_aux_loss_coef = kwargs.pop("moe_aux_loss_coef", 0.01)
         self._moe_metrics: Dict[str, float] = {}
-        self.extra_kwargs = kwargs
+        self.strategy_kwargs = kwargs
         self._rollout_runner = None
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:

astrai/trainer/train_context.py

@@ -140,8 +140,10 @@ class TrainContextBuilder:
                         checkpoint.consumed_samples // per_step * per_step
                     )
                     state.checkpoint = checkpoint
-        if not state.model_config and hasattr(cfg.model_fn(), "config"):
-            state.model_config = cfg.model_fn().config.to_dict()
+        if not state.model_config:
+            model = cfg.model_fn()
+            if hasattr(model, "config"):
+                state.model_config = model.config.to_dict()
         return state
 
     def _create_context(
@@ -260,7 +262,7 @@ class TrainContextBuilder:
 
     def _create_strategy(self, context: TrainContext, executor: BaseExecutor) -> dict:
         cfg = self.config
-        kwargs = dict(cfg.extra_kwargs)
+        kwargs = dict(cfg.strategy_kwargs)
         kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
         if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo"):
             kwargs["ref_model"] = create_ref_model(