1 files changed, 11 insertions, 2 deletions

diff --git a/data/keywords.py b/data/keywords.py
index 629006d..8632d67 100644
--- a/data/keywords.py
+++ b/data/keywords.py
@@ -1,14 +1,23 @@
+from typing import Optional
+
 import numpy as np
 
 
-def keywords_to_str(keywords: list[str], undroppable_keywords: list[str] = [], dropout: float = 0, shuffle: bool = False) -> str:
+def keywords_to_str(
+    keywords: list[str],
+    undroppable_keywords: list[str] = [],
+    dropout: float = 0,
+    shuffle: bool = False,
+    npgenerator: Optional[np.random.Generator] = None
+) -> str:
     if dropout != 0:
         keywords = [keyword for keyword in keywords if np.random.random() > dropout]
     else:
         keywords = keywords.copy()
     keywords += undroppable_keywords
     if shuffle:
-        np.random.shuffle(keywords)
+        npgenerator = npgenerator or np.random.default_rng()
+        npgenerator.shuffle(keywords)
     return ", ".join(keywords)