PR review

Author: jackzhxng

include/pytorch/tokenizers/re2_regex.h

@@ -1,12 +1,22 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #pragma once
 
 #include <memory>
 #include <string>
-#include "regex.h"
 
-// Third Party
 #include <re2/re2.h>
 
+#include <pytorch/tokenizers/regex.h>
+
+namespace tokenizers {
+
 /**
  * @brief RE2-based implementation of IRegex.
  */
@@ -24,11 +34,10 @@ class Re2Regex : public IRegex {
    */
   virtual std::vector<Match> findAll(const std::string& text) const override;
 
- protected:
   /**
    * @brief Check if RE2 compiled the pattern successfully.
    */
-  bool ok() const;
+  bool ok() const override;
 
   /**
    * @brief Expose internal RE2 pointer to the factory if needed.
@@ -38,5 +47,8 @@ class Re2Regex : public IRegex {
  private:
   std::unique_ptr<re2::RE2> regex_;
 
-  friend std::unique_ptr<IRegex> createRegex(const std::string& pattern);
+  friend Result<std::unique_ptr<IRegex>> createRegex(
+      const std::string& pattern);
 };
+
+} // namespace tokenizers

include/pytorch/tokenizers/regex.h

@@ -1,9 +1,21 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #pragma once
 
 #include <memory>
 #include <string>
 #include <vector>
 
+#include <pytorch/tokenizers/result.h>
+
+namespace tokenizers {
+
 struct Match {
   std::string text;
   size_t position;
@@ -23,6 +35,13 @@ class IRegex {
    * @return A vector of strings containing all matched substrings.
    */
   virtual std::vector<Match> findAll(const std::string& text) const = 0;
+
+  /**
+   * @brief Check if the regex pattern was compiled successfully.
+   *
+   * @return true if the pattern is valid and ready to use, false otherwise.
+   */
+  virtual bool ok() const = 0;
 };
 
 /**
@@ -31,4 +50,6 @@ class IRegex {
  * @param pattern The regex pattern to compile.
  * @return A unique pointer to an IRegex-compatible object.
  */
-std::unique_ptr<IRegex> createRegex(const std::string& pattern);
+Result<std::unique_ptr<IRegex>> createRegex(const std::string& pattern);
+
+} // namespace tokenizers

include/pytorch/tokenizers/std_regex.h

@@ -1,10 +1,20 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #pragma once
 
 #include <memory>
 #include <regex>
 #include <string>
 #include "regex.h"
 
+namespace tokenizers {
+
 /**
  * @brief std::regex-based implementation of IRegex.
  */
@@ -23,6 +33,15 @@ class StdRegex : public IRegex {
    */
   virtual std::vector<Match> findAll(const std::string& text) const override;
 
+  /**
+   * @brief Check if std::regex compiled the pattern successfully.
+   *
+   * @return true if the pattern is valid, false otherwise.
+   */
+  bool ok() const override;
+
  private:
   std::regex regex_;
 };
+
+} // namespace tokenizers

src/re2_regex.cpp

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <pytorch/tokenizers/re2_regex.h>
+
+namespace tokenizers {
+
+Re2Regex::Re2Regex(const std::string& pattern) {
+  regex_ = std::make_unique<re2::RE2>(pattern);
+  // Warmup re2 as it is slow on the first run, void the return value as it's
+  // not needed Refer to
+  // https://github.com/google/re2/blob/6dcd83d60f7944926bfd308cc13979fc53dd69ca/re2/fuzzing/re2_fuzzer.cc#L136-L141
+  (void)regex_->ReverseProgramSize();
+}
+
+std::vector<Match> Re2Regex::findAll(const std::string& text) const {
+  std::vector<Match> result;
+  re2::StringPiece input(text);
+  re2::StringPiece piece;
+
+  const char* base = input.data();
+
+  while (RE2::FindAndConsume(&input, *regex_, &piece)) {
+    size_t pos = piece.data() - base;
+    result.push_back({std::string(piece.data(), piece.size()), pos});
+  }
+
+  return result;
+}
+
+bool Re2Regex::ok() const {
+  return regex_ && regex_->ok();
+}
+
+const re2::RE2* Re2Regex::rawRegex() const {
+  return regex_.get();
+}
+
+} // namespace tokenizers

src/re2_regex.cpp

@@ -1,21 +1,23 @@
-#include "pytorch/tokenizers/regex.h"
-#include "pytorch/tokenizers/re2_regex.h"
-#include "pytorch/tokenizers/std_regex.h"
+#include <pytorch/tokenizers/regex.h>
+#include <pytorch/tokenizers/re2_regex.h>
+#include <pytorch/tokenizers/std_regex.h>
 
 #include <re2/re2.h>
 #include <iostream>
 #include <memory>
 
+namespace tokenizers {
+
 /**
  * @brief Factory function that creates a regex object using RE2 if possible.
  *        Falls back to std::regex if RE2 rejects the pattern with
- * ErrorBadPerlOp.
+ *        ErrorBadPerlOp.
  */
-std::unique_ptr<IRegex> createRegex(const std::string& pattern) {
-  auto re2 = std::make_unique<Re2Regex>(pattern);
+Result<std::unique_ptr<IRegex>> createRegex(const std::string& pattern) {
+  auto re2 = std::make_unique<Re2Regex>("(" + pattern + ")");
 
   if (re2->ok()) {
-    return re2;
+    return static_cast<std::unique_ptr<IRegex>>(std::move(re2));
   }
 
   const re2::RE2* raw = re2->rawRegex();
@@ -24,14 +26,17 @@ std::unique_ptr<IRegex> createRegex(const std::string& pattern) {
       std::cout
           << "RE2 is unable to support things such as negative lookaheads in "
           << pattern << ", defaulting to std::regex.";
-      return std::make_unique<StdRegex>(pattern);
+      auto std_regex = std::make_unique<StdRegex>("(" + pattern + ")");
+      return static_cast<std::unique_ptr<IRegex>>(std::move(std_regex));
     } catch (const std::regex_error& e) {
       std::cerr << "std::regex failed: " << e.what() << std::endl;
-      return nullptr;
+      return tokenizers::Error::LoadFailure;
     }
   } else {
     std::cerr << "RE2 failed to compile pattern: " << pattern << "\n";
     std::cerr << "Error: " << (raw ? raw->error() : "unknown") << std::endl;
-    return nullptr;
+    return tokenizers::Error::LoadFailure;
   }
 }
+
+} // namespace tokenizers

src/regex.cpp

@@ -1,9 +1,17 @@
-#include "pytorch/tokenizers/std_regex.h"
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <pytorch/tokenizers/std_regex.h>
 #include <regex>
 
-StdRegex::StdRegex(const std::string& pattern)
-    : regex_("(" + pattern + ")") // Add parentheses like RE2 version
-{}
+namespace tokenizers {
+
+StdRegex::StdRegex(const std::string& pattern) : regex_(pattern) {}
 
 std::vector<Match> StdRegex::findAll(const std::string& text) const {
   std::vector<Match> result;
@@ -20,3 +28,11 @@ std::vector<Match> StdRegex::findAll(const std::string& text) const {
 
   return result;
 }
+
+bool StdRegex::ok() const {
+  // std::regex constructor throws if the pattern is invalid
+  // If we got here, the pattern is valid
+  return true;
+}
+
+} // namespace tokenizers