uses abseil random library

Author: taku910

.gitignore

@@ -75,3 +75,5 @@ _sentencepiece.*.so
 third_party/abseil-cpp
 
 python/sentencepiece
+
+third_party/absl

data/gen_spec_parser.pl

@@ -18,6 +18,7 @@ list(APPEND SPM_LIBS absl::flags_parse)
 list(APPEND SPM_LIBS absl::log)
 list(APPEND SPM_LIBS absl::log_initialize)
 list(APPEND SPM_LIBS absl::check)
+list(APPEND SPM_LIBS absl::random_random)
 
 if (SPM_PROTOBUF_PROVIDER STREQUAL "internal")
   set(SPM_PROTO_HDRS builtin_pb/sentencepiece.pb.h)

src/CMakeLists.txt

@@ -23,6 +23,7 @@
 
 #include "freelist.h"
 #include "third_party/absl/container/flat_hash_map.h"
+#include "third_party/absl/random/random.h"
 #include "util.h"
 
 namespace sentencepiece {
@@ -129,7 +130,7 @@ std::vector<std::pair<absl::string_view, int>> Model::SampleEncode(
   }
 
   // BPE-dropout: https://arxiv.org/pdf/1910.13267.pdf
-  std::mt19937 *rand_gen = nullptr;
+  absl::BitGen *rand_gen = nullptr;
   auto skip_merge = [&]() {
     if (alpha <= 0.0) return false;
     if (alpha >= 1.0) return true;

src/bpe_model.cc

@@ -316,8 +316,7 @@ util::Status SentencePieceNormalizer::Load(absl::string_view filename) {
 util::Status SentencePieceNormalizer::LoadFromSerializedProto(
     absl::string_view serialized) {
   auto model_proto = std::make_unique<ModelProto>();
-  RET_CHECK(
-      model_proto->ParseFromArray(serialized.data(), serialized.size()));
+  RET_CHECK(model_proto->ParseFromArray(serialized.data(), serialized.size()));
   return Load(std::move(model_proto));
 }
 

src/sentencepiece_trainer.cc

@@ -30,6 +30,7 @@
 #include "sentencepiece_processor.h"
 #include "sentencepiece_trainer.h"
 #include "third_party/absl/container/flat_hash_map.h"
+#include "third_party/absl/random/random.h"
 #include "third_party/absl/strings/numbers.h"
 #include "third_party/absl/strings/str_cat.h"
 #include "third_party/absl/strings/str_format.h"
@@ -77,7 +78,7 @@ util::Status VerifySpec(const TrainerSpec &trainer_spec) {
 #undef CHECK_RANGE
 
   RET_CHECK(trainer_spec.input_sentence_size() <= 0 ||
-                  trainer_spec.input_sentence_size() > 100);
+            trainer_spec.input_sentence_size() > 100);
 
   RET_CHECK(!trainer_spec.unk_piece().empty());
   RET_CHECK(!trainer_spec.bos_piece().empty());
@@ -87,7 +88,7 @@ util::Status VerifySpec(const TrainerSpec &trainer_spec) {
   if (SentencePieceTrainer::GetPretokenizerForTraining() ||
       !trainer_spec.pretokenization_delimiter().empty()) {
     RET_CHECK(trainer_spec.model_type() == TrainerSpec::UNIGRAM ||
-                    trainer_spec.model_type() == TrainerSpec::BPE)
+              trainer_spec.model_type() == TrainerSpec::BPE)
         << "PretokenizerForTraining is only supported in UNIGRAM or BPE mode.";
   }
 
@@ -307,7 +308,7 @@ bool TrainerInterface::IsValidSentencePiece(
 }
 
 template <typename T>
-void AddDPNoise(const TrainerSpec &trainer_spec, std::mt19937 *generator,
+void AddDPNoise(const TrainerSpec &trainer_spec, absl::BitGen *generator,
                 T *to_update) {
   if (trainer_spec.differential_privacy_noise_level() > 0) {
     std::normal_distribution<float> dist(
@@ -327,13 +328,12 @@ util::Status TrainerInterface::LoadSentences() {
   RET_CHECK(sentences_.empty());
   RET_CHECK(required_chars_.empty());
   RET_CHECK(trainer_spec_.input_format().empty() ||
-                  trainer_spec_.input_format() == "text" ||
-                  trainer_spec_.input_format() == "tsv")
+            trainer_spec_.input_format() == "text" ||
+            trainer_spec_.input_format() == "tsv")
       << "Supported formats are 'text' and 'tsv'.";
 
-  RET_CHECK(
-      (sentence_iterator_ != nullptr && trainer_spec_.input().empty()) ||
-      (sentence_iterator_ == nullptr && !trainer_spec_.input().empty()))
+  RET_CHECK((sentence_iterator_ != nullptr && trainer_spec_.input().empty()) ||
+            (sentence_iterator_ == nullptr && !trainer_spec_.input().empty()))
       << "SentenceIterator and trainer_spec.input() must be exclusive.";
 
   RET_CHECK(
@@ -487,7 +487,7 @@ util::Status TrainerInterface::LoadSentences() {
           auto *generator = random::GetRandomGenerator();
           for (size_t i = n; i < sentences_.size(); i += num_workers) {
             AddDPNoise<int64_t>(trainer_spec_, generator,
-                              &(sentences_[i].second));
+                                &(sentences_[i].second));
           }
         });
       }
@@ -581,9 +581,8 @@ util::Status TrainerInterface::LoadSentences() {
 
   if (trainer_spec_.model_type() != TrainerSpec::WORD &&
       trainer_spec_.model_type() != TrainerSpec::CHAR) {
-    RET_CHECK_LE(
-        static_cast<int>(required_chars_.size() + meta_pieces_.size()),
-        trainer_spec_.vocab_size())
+    RET_CHECK_LE(static_cast<int>(required_chars_.size() + meta_pieces_.size()),
+                 trainer_spec_.vocab_size())
         << "Vocabulary size is smaller than required_chars. "
         << trainer_spec_.vocab_size() << " vs "
         << required_chars_.size() + meta_pieces_.size() << ". "
@@ -619,7 +618,7 @@ util::Status TrainerInterface::Serialize(ModelProto *model_proto) const {
 
   model_proto->Clear();
 
-#define CHECK_PIECE(piece)                                  \
+#define CHECK_PIECE(piece)                            \
   RET_CHECK(string_util::IsStructurallyValid(piece)); \
   RET_CHECK(!piece.empty());                          \
   RET_CHECK(dup.insert(piece).second) << piece << " is already defined";
@@ -656,17 +655,15 @@ util::Status TrainerInterface::Serialize(ModelProto *model_proto) const {
   if (!trainer_spec_.hard_vocab_limit() ||
       trainer_spec_.model_type() == TrainerSpec::CHAR) {
     RET_CHECK_GE(trainer_spec_.vocab_size(), model_proto->pieces_size());
-    RET_CHECK_GE(trainer_spec_.vocab_size(),
-                       static_cast<int32_t>(dup.size()));
+    RET_CHECK_GE(trainer_spec_.vocab_size(), static_cast<int32_t>(dup.size()));
     model_proto->mutable_trainer_spec()->set_vocab_size(
         model_proto->pieces_size());
   } else {
     RET_CHECK_EQ(trainer_spec_.vocab_size(), model_proto->pieces_size())
         << absl::StrFormat(
                "Vocabulary size too high (%d). Please set it to a value <= %d.",
                trainer_spec_.vocab_size(), model_proto->pieces_size());
-    RET_CHECK_EQ(trainer_spec_.vocab_size(),
-                       static_cast<int32_t>(dup.size()));
+    RET_CHECK_EQ(trainer_spec_.vocab_size(), static_cast<int32_t>(dup.size()));
   }
 
   // Saves self-testing data.

src/trainer_interface.cc

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.!
 
-#include <unordered_map>
-
 #include "third_party/absl/container/flat_hash_map.h"
 #include "unicode_script.h"
 #include "unicode_script_map.h"

src/unicode_script.cc

@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "third_party/absl/container/flat_hash_map.h"
+#include "third_party/absl/random/random.h"
 #include "third_party/absl/strings/str_split.h"
 #include "third_party/absl/strings/string_view.h"
 #include "util.h"

src/unigram_model.cc

@@ -123,7 +123,7 @@ class BoundedPriorityQueue {
   }
 
   size_t max_capacity_;
-  std::unordered_map<std::string, int64_t> data_;
+  absl::flat_hash_map<std::string, int64_t> data_;
 };
 }  // namespace
 

src/unigram_model_trainer.cc

@@ -18,24 +18,18 @@
 #include <cstddef>
 #include <memory>
 
+#include "third_party/absl/random/random.h"
+
 namespace sentencepiece {
 
 namespace {
-constexpr uint32_t kDefaultSeed = static_cast<uint32_t>(-1);
+static constexpr uint32_t kDefaultSeed = static_cast<uint32_t>(-1);
 static std::atomic<uint32_t> g_seed = kDefaultSeed;
 }  // namespace
 
-void SetRandomGeneratorSeed(uint32_t seed) {
-  if (seed != kDefaultSeed) g_seed.store(seed);
-}
+void SetRandomGeneratorSeed(uint32_t seed) { g_seed.store(seed); }
 
-uint32_t GetRandomGeneratorSeed() {
-  try {
-    return g_seed == kDefaultSeed ? std::random_device{}() : g_seed.load();
-  } catch (...) {
-    return g_seed.load();
-  }
-}
+uint32_t GetRandomGeneratorSeed() { return g_seed.load(); }
 
 namespace {
 std::shared_ptr<const std::string> *GetSharedDataDir() {
@@ -174,18 +168,14 @@ std::string UnicodeTextToUTF8(const UnicodeText &utext) {
 }  // namespace string_util
 
 namespace random {
-std::mt19937 *GetRandomGenerator() {
+absl::BitGen *GetRandomGenerator() {
   // Thread-locals occupy stack space in every thread ever created by the
   // program, even if that thread never uses the thread-local variable.
-  //
-  // https://maskray.me/blog/2021-02-14-all-about-thread-local-storage
-  //
-  // sizeof(std::mt19937) is several kilobytes, so it is safer to put that on
-  // the heap, leaving only a pointer to it in thread-local storage.  This must
-  // be a unique_ptr, not a raw pointer, so that the generator is not leaked on
-  // thread exit.
   thread_local static auto mt =
-      std::make_unique<std::mt19937>(GetRandomGeneratorSeed());
+      GetRandomGeneratorSeed() == kDefaultSeed
+          ? std::make_unique<absl::BitGen>()
+          : std::make_unique<absl::BitGen>(
+                std::seed_seq{GetRandomGeneratorSeed()});
   return mt.get();
 }
 }  // namespace random