From 776213e99da4ec389575e797d93de8d8960fa010 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sun, 9 Apr 2023 16:21:52 +0200
Subject: Update

---
 train_ti.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

(limited to 'train_ti.py')

diff --git a/train_ti.py b/train_ti.py
index ca5b113..ebac302 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -63,6 +63,12 @@ def parse_args():
         default=None,
         help="The name of the current project.",
     )
+    parser.add_argument(
+        "--auto_cycles",
+        type=int,
+        default=1,
+        help="How many cycles to run automatically."
+    )
     parser.add_argument(
         "--placeholder_tokens",
         type=str,
@@ -869,10 +875,15 @@ def main():
             mid_point=args.lr_mid_point,
         )
 
-        continue_training = True
-        training_iter = 1
+        training_iter = 0
+
+        while True:
+            training_iter += 1
+            if training_iter > args.auto_cycles:
+                response = input("Run another cycle? [y/n] ")
+                if response.lower().strip() == "n":
+                    break
 
-        while continue_training:
             print("")
             print(f"------------ TI cycle {training_iter} ------------")
             print("")
-- 
cgit v1.2.3-54-g00ecf