Merge pull request #1241 from google/internal-merge

Merges internal changes to OSS.

Author: taku910

src/bpe_model.cc

@@ -14,6 +14,7 @@
 
 #include "bpe_model.h"
 
+#include <cstdint>
 #include <functional>
 #include <memory>
 #include <queue>
@@ -22,8 +23,12 @@
 #include <vector>
 
 #include "freelist.h"
+#include "model_interface.h"
+#include "sentencepiece_model.pb.h"
+#include "third_party/absl/base/attributes.h"
 #include "third_party/absl/container/flat_hash_map.h"
 #include "third_party/absl/random/random.h"
+#include "third_party/absl/strings/string_view.h"
 #include "util.h"
 
 namespace sentencepiece {
@@ -43,17 +48,52 @@ std::vector<std::pair<absl::string_view, int>> Model::SampleEncode(
   }
 
   struct SymbolPair {
-    int left;     // left index of this pair
-    int right;    // right index of this pair
-    float score;  // score of this pair. large is better.
-    size_t size;  // length of this piece
+    union {
+      float score;  // score of this pair. large is better.
+      int32_t int_score;
+    };
+    uint32_t left;      // left index of this pair
+    int right;          // right index of this pair
+    unsigned int size;  // length of this piece
   };
 
   class SymbolPairComparator {
    public:
-    const bool operator()(SymbolPair *h1, SymbolPair *h2) {
-      return (h1->score < h2->score ||
-              (h1->score == h2->score && h1->left > h2->left));
+    ABSL_ATTRIBUTE_ALWAYS_INLINE inline bool operator()(const SymbolPair &h1,
+                                                        const SymbolPair &h2) {
+      const int32_t i1 = h1.int_score;
+      const int32_t i2 = h2.int_score;
+
+      // Fast path for the common case where both scores are negative because
+      // they are log-probabilities.
+      // Note: we use the fact that IEEE 754 floating point format enables
+      // to compare the integer representation of negative floats which is
+      // cheaper than using float comparison. And it works the same way for
+      // little endian and big endian machines because the IEEE 754 format is
+      // aligned with the endianness.
+      // `(i1 & i2) < 0` is an efficient way to check `i1 < 0 && i2 < 0`.
+      if ((i1 & i2) < 0) {
+        // For negative floats, their integer representation order is the
+        // reverse of the float order. That is, for two negative floats f1, f2,
+        // f1 < f2 iff i1 > i2.
+        return (i1 > i2) || (i1 == i2 && h1.left > h2.left);
+      }
+
+      // Slow path for uncommon cases (mixed signs or both positive).
+      // Note: the comparison between NaN and +0 and +1 can be different than
+      // if we used float numbers but it should not influence the result.
+      bool score_less;
+      // If signs are different ((i1 ^ i2) < 0), the negative score is smaller.
+      if ((i1 ^ i2) < 0) {
+        score_less = i1 < 0;
+      } else {
+        // If signs are the same (and not both negative), they must both be
+        // non-negative. For non-negative floats, integer order is the same as
+        // float order.
+        score_less = i1 < i2;
+      }
+
+      return score_less || (i1 == i2 && h1.left > h2.left);
     }
   };
 
@@ -64,50 +104,12 @@ std::vector<std::pair<absl::string_view, int>> Model::SampleEncode(
     absl::string_view piece;
   };
 
-  using Agenda = std::priority_queue<SymbolPair *, std::vector<SymbolPair *>,
-                                     SymbolPairComparator>;
-  Agenda agenda;
   std::vector<Symbol> symbols;
   symbols.reserve(normalized.size());
 
-  // Reverse merge rules.
-  // key: merged symbol, value: pair of original symbols.
-  absl::flat_hash_map<absl::string_view,
-                      std::pair<absl::string_view, absl::string_view>>
-      rev_merge;
-
-  // Pre-allocates SymbolPair for efficiency.
-  constexpr size_t kPreallocateSymbolPairSize = 256;
-  model::FreeList<SymbolPair> symbol_pair_allocator(kPreallocateSymbolPairSize);
-
-  // Lookup new symbol pair at [left, right] and inserts it to agenda.
-  auto MaybeAddNewSymbolPair = [this, &symbol_pair_allocator, &symbols, &agenda,
-                                &rev_merge](int left, int right) {
-    if (left == -1 || right == -1 || symbols[left].freeze ||
-        symbols[right].freeze)
-      return;
-    const absl::string_view piece(
-        symbols[left].piece.data(),
-        symbols[left].piece.size() + symbols[right].piece.size());
-    const auto it = pieces_.find(piece);
-    if (it == pieces_.end()) {
-      return;
-    }
-    auto *h = symbol_pair_allocator.Allocate();
-    h->left = left;
-    h->right = right;
-    h->score = GetScore(it->second);
-    h->size = piece.size();
-    agenda.push(h);
-
-    // Makes `rev_merge` for resegmentation.
-    if (IsUnusedInlined(it->second)) {
-      rev_merge[piece] =
-          std::make_pair(symbols[left].piece, symbols[right].piece);
-    }
-  };
-
-  // Splits the input into character sequence
+  // Splits the input into Symbols doing longest prefix match of the input
+  // from pieces(type:UNUSED) in the vocabulary.
+  // Does character splitting as a fallback of longest prefix match.
   int index = 0;
   while (!normalized.empty()) {
     Symbol s;
@@ -124,53 +126,118 @@ std::vector<std::pair<absl::string_view, int>> Model::SampleEncode(
     return {};
   }
 
+  std::vector<SymbolPair> agenda_vec;
+  agenda_vec.reserve(symbols.size());
+
+  // Reverse merge rules.
+  // key: merged symbol, value: pair of original symbols.
+  absl::flat_hash_map<absl::string_view,
+                      std::pair<absl::string_view, absl::string_view>>
+      rev_merge;
+
   // Lookup all bigrams.
-  for (size_t i = 1; i < symbols.size(); ++i) {
-    MaybeAddNewSymbolPair(i - 1, i);
+  if (symbols.size() > 1) {
+    int left = 0;
+    int right = 1;
+    Symbol *symbol_left = &symbols[left];
+    Symbol *symbol_right = &symbols[right];
+    for (; right < symbols.size();
+         left = right, symbol_left = symbol_right, ++right, ++symbol_right) {
+      if (symbol_left->freeze || symbol_right->freeze) continue;
+      const absl::string_view piece(
+          symbol_left->piece.data(),
+          symbol_left->piece.size() + symbol_right->piece.size());
+      const auto it = pieces_.find(piece);
+      if (it == pieces_.end()) continue;
+      SymbolPair &h = agenda_vec.emplace_back();
+      h.left = left;
+      h.right = right;
+      h.score = GetScore(it->second);
+      h.size = piece.size();
+
+      // Makes `rev_merge` for resegmentation.
+      if (IsUnusedInlined(it->second))
+        rev_merge[piece] =
+            std::make_pair(symbol_left->piece, symbol_right->piece);
+    }
   }
 
-  // BPE-dropout: https://arxiv.org/pdf/1910.13267.pdf
-  absl::BitGen *rand_gen = nullptr;
-  auto skip_merge = [&]() {
-    if (alpha <= 0.0) return false;
-    if (alpha >= 1.0) return true;
-    if (rand_gen == nullptr) rand_gen = random::GetRandomGenerator();
-    std::uniform_real_distribution<> gen(0.0, 1.0);
-    return gen(*rand_gen) < alpha;
+  using Agenda = std::priority_queue<SymbolPair, std::vector<SymbolPair>,
+                                     SymbolPairComparator>;
+  Agenda agenda(SymbolPairComparator(), std::move(agenda_vec));
+  // Lookup new symbol pair at [left, right] and inserts it to agenda.
+  auto MaybeAddNewSymbolPair = [this, &symbols, &agenda, &rev_merge](
+                                   int left, int right) {
+    if (left == -1 || right == -1) return;
+    const Symbol &left_symbol = symbols[left];
+    const Symbol &right_symbol = symbols[right];
+    if (left_symbol.freeze || right_symbol.freeze) return;
+    const absl::string_view piece(
+        left_symbol.piece.data(),
+        left_symbol.piece.size() + right_symbol.piece.size());
+    const auto it = pieces_.find(piece);
+    if (it == pieces_.end()) {
+      return;
+    }
+    const int id = it->second;
+    SymbolPair h;
+    h.left = left;
+    h.right = right;
+    h.score = GetScore(id);
+    h.size = piece.size();
+    agenda.push(h);
+
+    // Makes `rev_merge` for resegmentation.
+    if (IsUnusedInlined(id))
+      rev_merge[piece] = std::make_pair(left_symbol.piece, right_symbol.piece);
   };
 
+  absl::BitGen *rand_gen = nullptr;
   // Main loop.
   while (!agenda.empty()) {
-    SymbolPair *top = agenda.top();
-    agenda.pop();
-
-    // `top` is no longer available.
-    if (symbols[top->left].piece.empty() || symbols[top->right].piece.empty() ||
-        symbols[top->left].piece.size() + symbols[top->right].piece.size() !=
-            top->size) {
+    // Pop the top pair if it is stale.
+    const SymbolPair &top_ref = agenda.top();
+    if (symbols[top_ref.left].piece.empty() ||
+        symbols[top_ref.right].piece.empty() ||
+        (symbols[top_ref.left].piece.size() +
+             symbols[top_ref.right].piece.size() !=
+         top_ref.size)) {
+      agenda.pop();
       continue;
     }
 
-    // Note that orignal BPE-dropout paper assumes that all merged symbols are
-    // pre computed, but here we randomly skip merge opration inside this loop.
-    // This implemenation is theoretically equivalent to the original one.
-    if (skip_merge()) continue;
+    SymbolPair top = agenda.top();
+    agenda.pop();
+
+    Symbol &left_symbol = symbols[top.left];
+    Symbol &right_symbol = symbols[top.right];
+
+    // Note that original BPE-dropout paper assumes that all merged symbols are
+    // pre computed, but here we randomly skip merge operation inside this loop.
+    // This implementation is theoretically equivalent to the original one.
+    // BPE-dropout: https://arxiv.org/pdf/1910.13267.pdf
+    if (alpha > 0.0) {
+      if (alpha >= 1.0) continue;
+      if (rand_gen == nullptr) rand_gen = random::GetRandomGenerator();
+      std::uniform_real_distribution<> gen(0.0, 1.0);
+      if (gen(*rand_gen) < alpha) continue;
+    }
 
     // Replaces symbols with `top` rule.
-    symbols[top->left].piece = absl::string_view(
-        symbols[top->left].piece.data(),
-        symbols[top->left].piece.size() + symbols[top->right].piece.size());
+    left_symbol.piece =
+        absl::string_view(left_symbol.piece.data(),
+                          left_symbol.piece.size() + right_symbol.piece.size());
 
     // Updates prev/next pointers.
-    symbols[top->left].next = symbols[top->right].next;
-    if (symbols[top->right].next >= 0) {
-      symbols[symbols[top->right].next].prev = top->left;
+    left_symbol.next = right_symbol.next;
+    if (right_symbol.next >= 0) {
+      symbols[right_symbol.next].prev = top.left;
     }
-    symbols[top->right].piece = absl::string_view("");
+    right_symbol.piece = absl::string_view("");
 
     // Adds new symbol pairs which are newly added after symbol replacement.
-    MaybeAddNewSymbolPair(symbols[top->left].prev, top->left);
-    MaybeAddNewSymbolPair(top->left, symbols[top->left].next);
+    MaybeAddNewSymbolPair(left_symbol.prev, top.left);
+    MaybeAddNewSymbolPair(top.left, left_symbol.next);
   }
 
   std::function<void(absl::string_view, EncodeResult *)> resegment;
@@ -194,6 +261,7 @@ std::vector<std::pair<absl::string_view, int>> Model::SampleEncode(
   };
 
   EncodeResult output;
+  output.reserve(symbols.size());
   for (int index = 0; index != -1; index = symbols[index].next) {
     if (index >= 0 && index < static_cast<int>(symbols.size())) {
       resegment(symbols[index].piece, &output);