fix: make sampling more reproducible

huggingface · Nov 23, 2023 · 2997147 · 2997147
1 parent 49cc545
commit 2997147
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/setfit/data.py b/src/setfit/data.py
@@ -170,7 +170,7 @@ def sample_dataset(dataset: Dataset, label_column: str = "label", num_samples: i
     df = df.groupby(label_column)
 
     # sample num_samples, or at least as much as possible
-    df = df.apply(lambda x: x.sample(min(num_samples, len(x))))
+    df = df.apply(lambda x: x.sample(min(num_samples, len(x)), random_state=seed))
     df = df.reset_index(drop=True)
 
     all_samples = Dataset.from_pandas(df, features=dataset.features)