only extra-subsample legal

segment-any-text · May 14, 2024 · e6a88f8 · e6a88f8
1 parent 2e80a9d
commit e6a88f8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/wtpsplit/train/train_adapter.py b/wtpsplit/train/train_adapter.py
@@ -127,7 +127,7 @@ def prepare_dataset(
                         processed_chunk[args.text_column] = "\n".join(chunk)
                         processed_dataset.append(processed_chunk)
                     dataset = datasets.Dataset.from_list(processed_dataset)
-                    if subsample:
+                    if subsample and "legal" in dataset_name:
                         # 10k sentences -> 1k documents.
                         subsample = subsample // 10