don't store logits if no punct applied

segment-any-text · May 15, 2024 · a4f91af · a4f91af
1 parent 1825000
commit a4f91af
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/wtpsplit/evaluation/intrinsic.py b/wtpsplit/evaluation/intrinsic.py
@@ -268,7 +268,13 @@ def load_or_compute_logits(args, model, eval_data, valid_data=None, save_str: st
                         dset_group.create_dataset("test_logit_lengths", data=test_logit_lengths)
                     else:
                         test_labels = get_labels(lang_code, test_sentences, after_space=False)
-
+                    if args.skip_punct:
+                        print(test_logits[:, 0].shape)
+                        # remove punct logits
+                        test_logits = test_logits[:, 0]
+                        # back to [N, 1]
+                        test_logits = np.expand_dims(test_logits, axis=1)
+                        print(test_logits.shape)
                     dset_group.create_dataset("test_logits", data=test_logits)
                     dset_group.create_dataset("test_labels", data=test_labels)
 
@@ -296,6 +302,11 @@ def load_or_compute_logits(args, model, eval_data, valid_data=None, save_str: st
                     else:
                         train_labels = get_labels(lang_code, train_sentences, after_space=False)
 
+                    if args.skip_punct:
+                        # remove punct logits
+                        train_logits = train_logits[:, 0]
+                        # back to [N, 1]
+                        train_logits = np.expand_dims(train_logits, axis=1)
                     dset_group.create_dataset("train_logits", data=train_logits)
                     dset_group.create_dataset("train_labels", data=train_labels)