Added sum1 functionality

configs/base.py

@@ -52,6 +52,12 @@ def get_config():
     training.ckpt_best_metric_best_mode = "min"  # "max" or "min" 
     training.ckpt_keep_top_n = 3  # Save the top `ckpt_keep_top_n` checkpoints based on `ckpt_best_metric`
 
+    # Sum1 options
+    config.sum1 = sum1 = ml_collections.ConfigDict()
+    sum1.use_sum1 = True
+    sum1.summaries_filename = "summary_v2"
+    sum1.sum1_filename = "summary_sum1_v2"
+
     # Optimizer (AdamW)
     config.optim = optim = ml_collections.ConfigDict()
     optim.schedule = "constant"

environment.yml

@@ -35,7 +35,6 @@ dependencies:
   - six=1.16.0=pyh6c4a22f_0
   - tk=8.6.13=noxft_h4845f30_101
   - traitlets=5.14.0=pyhd8ed1ab_0
-  - tzdata=2023d=h0c530f3_0
   - wheel=0.42.0=pyhd8ed1ab_0
   - xz=5.2.6=h166bdaf_0
   - zeromq=4.3.5=h59595ed_0
@@ -163,6 +162,7 @@ dependencies:
       - orbax-checkpoint==0.4.8
       - overrides==7.4.0
       - packaging==23.2
+      - pandas==2.1.4
       - pandocfilters==1.5.0
       - parso==0.8.3
       - pexpect==4.9.0
@@ -180,6 +180,7 @@ dependencies:
       - pygments==2.17.2
       - pyparsing==3.1.1
       - python-json-logger==2.0.7
+      - pytz==2023.3.post1
       - pyvo==1.5
       - pyyaml==6.0.1
       - qtconsole==5.5.1
@@ -195,6 +196,7 @@ dependencies:
       - rsa==4.9
       - safetensors==0.4.1
       - scipy==1.11.4
+      - seaborn==0.13.1
       - secretstorage==3.3.3
       - send2trash==1.8.2
       - sentry-sdk==1.39.1
@@ -220,6 +222,7 @@ dependencies:
       - transformers==4.36.2
       - types-python-dateutil==2.8.19.14
       - typing-extensions==4.9.0
+      - tzdata==2023.4
       - uri-template==1.3.0
       - urllib3==2.1.0
       - wandb==0.16.1