updates

modeling/official/modeling/multitask/base_trainer.py

@@ -19,12 +19,12 @@ The trainer derives from the Orbit `StandardTrainer` class.
 from typing import Union
 
 import gin
-import orbit
+import modeling.orbit
 import tensorflow as tf, tf_keras
 
-from official.modeling import optimization
-from official.modeling.multitask import base_model
-from official.modeling.multitask import multitask
+from modeling.official.modeling import optimization
+from modeling.official.modeling.multitask import base_model
+from modeling.official.modeling.multitask import multitask
 
 
 @gin.configurable

modeling/official/modeling/multitask/configs.py

@@ -16,9 +16,9 @@
 import dataclasses
 from typing import Optional, Tuple
 
-from official.core import config_definitions as cfg
-from official.modeling import hyperparams
-from official.modeling.privacy import configs as dp_configs
+from modeling.official.core import config_definitions as cfg
+from modeling.official.modeling import hyperparams
+from modeling.official.modeling.privacy import configs as dp_configs
 
 
 @dataclasses.dataclass

modeling/official/modeling/multitask/evaluator.py

@@ -18,12 +18,12 @@ The evaluator implements the Orbit `AbstractEvaluator` interface.
 """
 from typing import Dict, List, Optional, Union
 import gin
-import orbit
+import modeling.orbit
 import tensorflow as tf, tf_keras
 
-from official.core import base_task
-from official.core import train_utils
-from official.modeling.multitask import base_model
+from modeling.official.core import base_task
+from modeling.official.core import train_utils
+from modeling.official.modeling.multitask import base_model
 
 
 @gin.configurable

modeling/official/modeling/multitask/interleaving_trainer.py

@@ -15,12 +15,12 @@
 """Multitask trainer that interleaves each task's train step."""
 from typing import Union
 import gin
-import orbit
+import modeling.orbit
 import tensorflow as tf, tf_keras
-from official.modeling.multitask import base_model
-from official.modeling.multitask import base_trainer
-from official.modeling.multitask import multitask
-from official.modeling.multitask import task_sampler as sampler
+from modeling.official.modeling.multitask import base_model
+from modeling.official.modeling.multitask import base_trainer
+from modeling.official.modeling.multitask import multitask
+from modeling.official.modeling.multitask import task_sampler as sampler
 
 
 @gin.configurable

modeling/official/modeling/multitask/multitask.py

@@ -17,13 +17,13 @@ import abc
 from typing import Dict, List, Optional, Text, Union
 
 import tensorflow as tf, tf_keras
-from official.core import base_task
-from official.core import config_definitions
-from official.core import task_factory
-from official.modeling import optimization
-from official.modeling.multitask import base_model
-from official.modeling.multitask import configs
-from official.modeling.privacy import configs as dp_configs
+from modeling.official.core import base_task
+from modeling.official.core import config_definitions
+from modeling.official.core import task_factory
+from modeling.official.modeling import optimization
+from modeling.official.modeling.multitask import base_model
+from modeling.official.modeling.multitask import configs
+from modeling.official.modeling.privacy import configs as dp_configs
 
 OptimizationConfig = optimization.OptimizationConfig
 RuntimeConfig = config_definitions.RuntimeConfig

modeling/official/modeling/multitask/task_sampler.py

@@ -17,7 +17,7 @@ import abc
 from typing import Union, Dict, Text
 import tensorflow as tf, tf_keras
 
-from official.modeling.multitask import configs
+from modeling.official.modeling.multitask import configs
 
 
 class TaskSampler(tf.Module, metaclass=abc.ABCMeta):

modeling/official/modeling/multitask/test_utils.py

@@ -15,10 +15,10 @@
 """Testing utils for mock models and tasks."""
 from typing import Dict, Text
 import tensorflow as tf, tf_keras
-from official.core import base_task
-from official.core import config_definitions as cfg
-from official.core import task_factory
-from official.modeling.multitask import base_model
+from modeling.official.core import base_task
+from modeling.official.core import config_definitions as cfg
+from modeling.official.core import task_factory
+from modeling.official.modeling.multitask import base_model
 
 
 class MockFooModel(tf_keras.Model):

modeling/official/modeling/multitask/train_lib.py

@@ -17,18 +17,18 @@
 import os
 from typing import Any, List, Mapping, Optional, Tuple, Union
 from absl import logging
-import orbit
+import modeling.orbit
 import tensorflow as tf, tf_keras
-from official.core import base_task
-from official.core import base_trainer as core_lib
-from official.core import train_utils
-from official.modeling.multitask import base_model
-from official.modeling.multitask import base_trainer
-from official.modeling.multitask import configs
-from official.modeling.multitask import evaluator as evaluator_lib
-from official.modeling.multitask import interleaving_trainer
-from official.modeling.multitask import multitask
-from official.modeling.multitask import task_sampler
+from modeling.official.core import base_task
+from modeling.official.core import base_trainer as core_lib
+from modeling.official.core import train_utils
+from modeling.official.modeling.multitask import base_model
+from modeling.official.modeling.multitask import base_trainer
+from modeling.official.modeling.multitask import configs
+from modeling.official.modeling.multitask import evaluator as evaluator_lib
+from modeling.official.modeling.multitask import interleaving_trainer
+from modeling.official.modeling.multitask import multitask
+from modeling.official.modeling.multitask import task_sampler
 
 TRAINERS = {
     'interleaving': interleaving_trainer.MultiTaskInterleavingTrainer,

modeling/official/modeling/privacy/configs.py

@@ -15,7 +15,7 @@
 """Configs for differential privacy."""
 import dataclasses
 
-from official.modeling.hyperparams import base_config
+from modeling.official.modeling.hyperparams import base_config
 
 
 @dataclasses.dataclass

modeling/official/modeling/privacy/configs_test.py

@@ -15,7 +15,7 @@
 """Tests for configs."""
 
 import tensorflow as tf, tf_keras
-from official.modeling.privacy import configs
+from modeling.official.modeling.privacy import configs
 
 
 class ConfigsTest(tf.test.TestCase):

modeling/official/modeling/privacy/ops_test.py

@@ -18,7 +18,7 @@ from unittest import mock
 
 import tensorflow as tf, tf_keras
 
-from official.modeling.privacy import ops
+from modeling.official.modeling.privacy import ops
 
 
 class OpsTest(tf.test.TestCase):