Minor style improvements.

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0  # Use the ref you want to point at
+    rev: v4.6.0  # Use the ref you want to point at
     hooks:
     - id: trailing-whitespace
     - id: check-ast
@@ -15,8 +15,8 @@ repos:
     - id: check-yaml
     - id: check-toml
 
-  - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.3.7'
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: 'v0.4.2'
     hooks:
     - id: ruff
       args: [--fix, --exit-non-zero-on-fix]
@@ -25,7 +25,7 @@ repos:
       types_or: [python, jupyter]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.8.0
+    rev: v1.10.0
     hooks:
     - id: mypy
       entry: python3 -m mypy --config-file pyproject.toml
@@ -41,3 +41,22 @@ repos:
       language: system
       pass_filenames: false
       always_run: true
+
+  - repo: local
+    hooks:
+    - id: nbstripout
+      name: nbstripout
+      language: system
+      entry: python3 -m nbstripout
+
+ci:
+    autofix_commit_msg: |
+        [pre-commit.ci] Add auto fixes from pre-commit.com hooks
+
+        for more information, see https://pre-commit.ci
+    autofix_prs: true
+    autoupdate_branch: ''
+    autoupdate_commit_msg: '[pre-commit.ci] pre-commit autoupdate'
+    autoupdate_schedule: weekly
+    skip: [pytest,nbstripout,mypy]
+    submodules: false

finetune.py

@@ -32,12 +32,11 @@
 
 
 def main(
-    args: Dict[str, Any],
+    args: argparse.Namespace,
     pre_model_config: Dict[str, Any],
     fine_model_config: Dict[str, Any],
 ) -> None:
     """Train the model."""
-
     # Setup environment
     seed_everything(args.seed)
     os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
@@ -66,7 +65,7 @@ def main(
                 random_state=args.seed,
                 stratify=fine_tune["label"],
             )
-        
+
         else:  # Multi label classfication
             fine_train_ids, _, fine_val_ids, _ = iterative_train_test_split(
                 X=fine_tune["patient_id"].to_numpy().reshape(-1, 1),
@@ -113,7 +112,7 @@ def main(
             balance_guide=None,
             max_len=args.max_len,
         )
-    
+
     else:
         train_dataset = FinetuneDataset(
             data=fine_train,
@@ -130,7 +129,7 @@ def main(
             tokenizer=tokenizer,
             max_len=args.max_len,
         )
-    
+
     train_loader = DataLoader(
         train_dataset,
         batch_size=args.batch_size,
@@ -163,12 +162,12 @@ def main(
             dirpath=args.checkpoint_dir,
         ),
         LearningRateMonitor(logging_interval="step"),
-        # EarlyStopping(
-        #     monitor="val_loss",
-        #     patience=args.patience,
-        #     verbose=True,
-        #     mode="min",
-        # ),
+        EarlyStopping(
+            monitor="val_loss",
+            patience=args.patience,
+            verbose=True,
+            mode="min",
+        ),
     ]
 
     # Create model
@@ -185,7 +184,7 @@ def main(
             pretrained_model=pretrained_model,
             **fine_model_config,
         )
-    
+
     elif args.model_type == "cehr_bigbird":
         pretrained_model = BigBirdPretrain(
             vocab_size=tokenizer.get_vocab_size(),

interpret.py

@@ -25,7 +25,7 @@
 
 
 def main(
-    args: Dict[str, Any],
+    args: argparse.Namespace,
     pre_model_config: Dict[str, Any],
     fine_model_config: Dict[str, Any],
 ) -> None: