Internal change

PiperOrigin-RevId: 330660757

official/core/train_utils.py

@@ -18,9 +18,10 @@
 import json
 import os
 import pprint
-from typing import Any
+from typing import Any, List
 
 from absl import logging
+import dataclasses
 import orbit
 import tensorflow as tf
 
@@ -47,6 +48,16 @@ def create_trainer(
   return trainer
 
 
+@dataclasses.dataclass
+class ParseConfigOptions:
+  """Use this dataclass instead of FLAGS to customize parse_configuration()."""
+  experiment: str
+  config_file: List[str]
+  tpu: str = ''
+  tf_data_service: str = ''
+  params_override: str = ''
+
+
 def parse_configuration(flags_obj):
   """Parses ExperimentConfig from flags."""