bug fix dummy task (put in schema)

AllenNeuralDynamics · Dec 20, 2023 · f287c96 · f287c96
1 parent 29e3e8f
commit f287c96
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 27 deletions.
diff --git a/code/aind_auto_train/curriculums/dummy_task.py b/code/aind_auto_train/curriculums/dummy_task.py
@@ -4,31 +4,8 @@
 from typing import List, Dict
 
 from aind_auto_train.curriculum_manager import LOCAL_SAVED_CURRICULUM_ROOT
-from aind_auto_train.schema.task import Task, TrainingStage, Metrics, TaskParas
-from aind_auto_train.schema.curriculum import Curriculum, StageTransitions, TransitionRule, Decision
-
-
-# Override the metrics class
-class DummyTaskMetrics(Metrics):
-    dummy_metric_float: List[float]
-    dummy_metric_int: List[int]
-
-# Override the task parameters class
-class DummyTaskParas(TaskParas):
-    dummy_para_bool: bool
-    dummy_para_float: float
-
-# Override the curriculum class
-class DummyTaskCurriculum(Curriculum[DummyTaskParas, DummyTaskMetrics]):
-    # Override parameters
-    parameters: Dict[TrainingStage, DummyTaskParas]
-
-    # Override metrics
-    def evaluate_transitions(self,
-                             current_stage: TrainingStage,
-                             metrics: DummyTaskMetrics  # Note the dynamical type here
-                             ) -> TrainingStage:
-        return super().evaluate_transitions(current_stage, metrics)
+from aind_auto_train.schema.task import Task, TrainingStage, DummyTaskParas, DummyTaskMetrics
+from aind_auto_train.schema.curriculum import DummyTaskCurriculum, StageTransitions, TransitionRule, Decision
 
 
 meta = dict(curriculum_version="0.1",

diff --git a/code/aind_auto_train/schema/curriculum.py b/code/aind_auto_train/schema/curriculum.py
@@ -10,8 +10,8 @@
 from pydantic.json import pydantic_encoder
 
 from aind_auto_train.schema.task import (Task, TrainingStage,
-                                         taskparas_class, DynamicForagingParas,
-                                         metrics_class, DynamicForagingMetrics)
+                                         taskparas_class, DynamicForagingParas, DummyTaskParas,
+                                         metrics_class, DynamicForagingMetrics, DummyTaskMetrics)
 from aind_auto_train.plot.curriculum import draw_diagram_rules, draw_diagram_paras
 
 # %%
@@ -192,6 +192,18 @@ def evaluate_transitions(self,
         return super().evaluate_transitions(current_stage, metrics)
 
 
+class DummyTaskCurriculum(Curriculum[DummyTaskParas, DummyTaskMetrics]):
+    # Override parameters
+    parameters: Dict[TrainingStage, DummyTaskParas]
+
+    # Override metrics
+    def evaluate_transitions(self,
+                             current_stage: TrainingStage,
+                             metrics: DummyTaskMetrics  # Note the dynamical type here
+                             ) -> TrainingStage:
+        return super().evaluate_transitions(current_stage, metrics)
+
+
 # ------------------ Helpers ------------------
 # A hack to serialize TrainingStage in the dictionary keys
 def transform_dict_with_enum_keys(obj):

diff --git a/code/aind_auto_train/schema/task.py b/code/aind_auto_train/schema/task.py
@@ -52,6 +52,11 @@ class DynamicForagingMetrics(Metrics):
     foraging_efficiency: List[float]  # Full history of foraging efficiency
     finished_trials: List[int]  # Full history of finished trials
 
+# For dummy task
+class DummyTaskMetrics(Metrics):
+    dummy_metric_float: List[float]
+    dummy_metric_int: List[int]
+
 class TaskParas(AindModel):
     """Parent class for TaskParas. All other task parameters should inherit from this class
     """
@@ -156,3 +161,8 @@ class DynamicForagingParas(TaskParas):
 
 
 
+# For dummy task
+class DummyTaskParas(TaskParas):
+    dummy_para_bool: bool
+    dummy_para_float: float
+