1 files changed, 5 insertions, 7 deletions

diff --git a/train_ti.py b/train_ti.py
index c1c0eed..48858cc 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -17,7 +17,6 @@ import transformers
 from util.files import load_config, load_embeddings_from_dir
 from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, add_placeholder_tokens, get_models
-from training.lr import plot_metrics
 from training.strategy.ti import textual_inversion_strategy
 from training.optimization import get_scheduler
 from training.util import save_args
@@ -511,12 +510,12 @@ def parse_args():
     if isinstance(args.initializer_tokens, str):
         args.initializer_tokens = [args.initializer_tokens] * len(args.placeholder_tokens)
 
-    if len(args.initializer_tokens) == 0:
-        raise ValueError("You must specify --initializer_tokens")
-
     if len(args.placeholder_tokens) == 0:
         args.placeholder_tokens = [f"<*{i}>" for i in range(len(args.initializer_tokens))]
 
+    if len(args.initializer_tokens) == 0:
+        args.initializer_tokens = args.placeholder_tokens.copy()
+
     if len(args.placeholder_tokens) != len(args.initializer_tokens):
         raise ValueError("--placeholder_tokens and --initializer_tokens must have the same number of items")
 
@@ -856,7 +855,7 @@ def main():
             mid_point=args.lr_mid_point,
         )
 
-        metrics = trainer(
+        trainer(
             project="textual_inversion",
             train_dataloader=datamodule.train_dataloader,
             val_dataloader=datamodule.val_dataloader,
@@ -864,14 +863,13 @@ def main():
             lr_scheduler=lr_scheduler,
             num_train_epochs=num_train_epochs,
             # --
+            group_labels=["emb"],
             sample_output_dir=sample_output_dir,
             sample_frequency=sample_frequency,
             placeholder_tokens=placeholder_tokens,
             placeholder_token_ids=placeholder_token_ids,
         )
 
-        plot_metrics(metrics, metrics_output_file)
-
     if not args.sequential:
         run(0, args.placeholder_tokens, args.initializer_tokens, args.num_vectors, args.train_data_template)
     else: