Adding Numpy random number generator

mc4/mc4.py

@@ -7,6 +7,8 @@ import json
 import datasets
 import kenlm
 import numpy as np
+from numpy.random import default_rng
+rng = default_rng()
 
 
 logger = datasets.logging.get_logger(__name__)
@@ -309,7 +311,6 @@ class Mc4(datasets.GeneratorBasedBuilder):
             doc_length += length
         return 10.0 ** (-doc_log_score / doc_length)
 
-
     def _should_keep_doc_step(self, doc, factor=1, boundaries=None):
         perplexity = self.get_perplexity(doc)
         if boundaries is None:
@@ -323,7 +324,7 @@ class Mc4(datasets.GeneratorBasedBuilder):
         elif perplexity >= boundaries[2]:
             quartile_range = 100 * boundaries[2]
         probability = factor / quartile_range
-        return np.random() < probability
+        return rng.uniform() < probability
 
     def _should_keep_doc_gaussian(self, doc, factor=0.4, boundaries=None):
         perplexity = self.get_perplexity(doc)
@@ -332,10 +333,10 @@ class Mc4(datasets.GeneratorBasedBuilder):
         else:
             m = 662247.50212365
         weighted_perplexity = factor * np.exp(-9/2*((perplexity-m)/m)**2)
-        return np.random.uniform() < weighted_perplexity
+        return rng.uniform() < weighted_perplexity
 
     def _should_keep_doc_random(self, doc, factor=None, boundaries=None):
-        return np.random() <= 0.5
+        return rng.uniform() <= 0.5
 
     def _info(self):
         return datasets.DatasetInfo(