Add python bindings

Author: larryliu0820

CMakeLists.txt

@@ -18,6 +18,7 @@ project(Tokenizers)
 
 option(TOKENIZERS_BUILD_TEST "Build tests" OFF)
 option(TOKENIZERS_BUILD_TOOLS "Build tools" OFF)
+option(TOKENIZERS_BUILD_PYTHON "Build Python bindings" OFF)
 option(SUPPORT_REGEX_LOOKAHEAD
        "Support regex lookahead patterns (requires PCRE2)" OFF
 )
@@ -121,6 +122,38 @@ if(TOKENIZERS_BUILD_TOOLS)
   add_subdirectory(examples/tokenize_tool)
 endif()
 
+# Build Python bindings
+if(TOKENIZERS_BUILD_PYTHON)
+  include(FetchContent)
+  FetchContent_Declare(
+    pybind11
+    GIT_REPOSITORY https://github.com/pybind/pybind11.git
+    GIT_TAG v2.13.6
+  )
+  FetchContent_MakeAvailable(pybind11)
+  # Create the Python extension module
+  pybind11_add_module(pytorch_tokenizers_cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/python_bindings.cpp
+  )
+
+  # Link with the tokenizers library
+  target_link_libraries(pytorch_tokenizers_cpp PRIVATE tokenizers)
+
+  # Set properties for the Python extension
+  target_compile_definitions(pytorch_tokenizers_cpp PRIVATE VERSION_INFO=${PROJECT_VERSION})
+
+  # Set the output name and location
+  set_target_properties(pytorch_tokenizers_cpp PROPERTIES
+    OUTPUT_NAME "pytorch_tokenizers_cpp"
+  )
+
+  # Install the Python extension
+  install(TARGETS pytorch_tokenizers_cpp
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+  )
+endif()
+
 # Installation rules
 include(GNUInstallDirs)
 

include/pytorch/tokenizers/tiktoken.h

@@ -46,6 +46,27 @@ class Tiktoken : public detail::BPETokenizerBase {
     }
   }
 
+  explicit Tiktoken(
+      std::string pattern,
+      const std::vector<std::string>& special_tokens,
+      size_t bos_token_index,
+      size_t eos_token_index)
+      : Tiktoken(
+            pattern,
+            std::make_unique<std::vector<std::string>>(special_tokens),
+            bos_token_index,
+            eos_token_index) {}
+
+  explicit Tiktoken(
+      const std::vector<std::string>& special_tokens,
+      size_t bos_token_index,
+      size_t eos_token_index)
+      : Tiktoken(
+            _get_default_patern(),
+            std::make_unique<std::vector<std::string>>(special_tokens),
+            bos_token_index,
+            eos_token_index) {}
+
   explicit Tiktoken(
       std::unique_ptr<std::vector<std::string>> special_tokens,
       size_t bos_token_index,

pyproject.toml

@@ -64,12 +64,22 @@ Changelog = "https://github.com/pytorch/executorch/releases"
 [tool.setuptools.exclude-package-data]
 "*" = ["*.pyc"]
 
-[tool.usort]
-# Do not try to put "first-party" imports in their own section.
-first_party_detection = false
+[tool.pytest.ini_options]
+testpaths = ["test"]
+python_files = ["test_*.py", "*_test.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
 
 [tool.black]
-# Emit syntax compatible with older versions of python instead of only the range
-# specified by `requires-python`. TODO: Remove this once we support these older
-# versions of python and can expand the `requires-python` range.
-target-version = ["py38", "py39", "py310", "py311", "py312"]
+target-version = ['py38', 'py39', 'py310', 'py311', 'py312']
+include = '\.pyi?$'
+extend-exclude = '''
+/(
+  # directories
+  \.eggs
+  | \.git
+  | build
+  | dist
+  | third-party
+)/
+'''

pytest.ini

@@ -0,0 +1,51 @@
+[pytest]
+# Pytest configuration for PyTorch Tokenizers
+
+# Test discovery - be very explicit about what to test
+testpaths = test
+python_files = test_*.py *_test.py
+python_classes = Test*
+python_functions = test_*
+
+# Add the project root to Python path
+pythonpath = .
+
+# Output options with explicit ignores
+addopts =
+    -v
+    --tb=short
+    --strict-markers
+    --disable-warnings
+    --ignore=third-party
+    --ignore=build
+    --ignore=cmake
+    --ignore=examples
+    --ignore=pytorch_tokenizers.egg-info
+
+# Directories to ignore during test collection
+norecursedirs =
+    build*
+    third-party*
+    cmake*
+    examples*
+    .git*
+    __pycache__*
+    *.egg-info*
+    *third-party*
+
+# Markers
+markers =
+    slow: marks tests as slow (deselect with '-m "not slow"')
+    integration: marks tests as integration tests
+    unit: marks tests as unit tests
+
+# Minimum version
+minversion = 6.0
+
+# Test timeout (in seconds)
+timeout = 300
+
+# Filter warnings
+filterwarnings =
+    ignore::DeprecationWarning
+    ignore::PendingDeprecationWarning

pytorch_tokenizers/__init__.py

@@ -3,25 +3,38 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-# @lint-ignore-every LICENSELINT
 
+"""
+PyTorch Tokenizers - Fast tokenizers for PyTorch
 
-from typing import Optional
+This package provides Python bindings for fast C++ tokenizer implementations
+including HuggingFace, TikToken, Llama2C, and SentencePiece tokenizers.
+"""
 
-from .hf_tokenizer import HuggingFaceTokenizer
-from .llama2c import Llama2cTokenizer
-from .tiktoken import TiktokenTokenizer
+try:
+    from .pytorch_tokenizers_cpp import (
+        Error,
+        TokenIndex,
+        Tokenizer,
+        HFTokenizer,
+        Tiktoken,
+        Llama2cTokenizer,
+        SPTokenizer,
+    )
+except ImportError as e:
+    raise ImportError(
+        f"Failed to import C++ tokenizer bindings: {e}. "
+        "Make sure the package was built correctly with pybind11."
+    ) from e
 
-__all__ = ["TiktokenTokenizer", "Llama2cTokenizer", "HuggingFaceTokenizer"]
+__version__ = "0.1.0"
 
-
-def get_tokenizer(tokenizer_path: str, tokenizer_config_path: Optional[str] = None):
-    if tokenizer_path.endswith(".json"):
-        tokenizer = HuggingFaceTokenizer(tokenizer_path, tokenizer_config_path)
-    else:
-        try:
-            tokenizer = Llama2cTokenizer(model_path=str(tokenizer_path))
-        except Exception:
-            print("Using Tiktokenizer")
-            tokenizer = TiktokenTokenizer(model_path=str(tokenizer_path))
-    return tokenizer
+__all__ = [
+    "Error",
+    "TokenIndex",
+    "Tokenizer",
+    "HFTokenizer",
+    "Tiktoken",
+    "Llama2cTokenizer",
+    "SPTokenizer",
+]

requirements.txt

@@ -0,0 +1,10 @@
+# Build dependencies
+pybind11>=2.6.0
+setuptools>=45
+wheel
+cmake>=3.18
+ninja
+
+# Optional dependencies for development
+pytest>=6.0
+numpy>=1.19.0

setup.py

@@ -5,14 +5,151 @@
 # LICENSE file in the root directory of this source tree.
 # @lint-ignore-every LICENSELINT
 # type: ignore[syntax]
-from setuptools import find_packages, setup
 
+import os
+import re
+import subprocess
+import sys
+import shutil
+from pathlib import Path
+
+from setuptools import find_packages, setup, Extension
+from setuptools.command.build_ext import build_ext
+
+# Read the README file
 with open("README.md", "r") as f:
     long_description = f.read()
 
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=""):
+        Extension.__init__(self, name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext):
+        extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
+
+        # Required for auto-detection & inclusion of auxiliary "native" libs
+        if not extdir.endswith(os.path.sep):
+            extdir += os.path.sep
+
+        debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
+        cfg = "Debug" if debug else "Release"
+
+        # CMake lets you override the generator - we check this.
+        # Can be set with Conda-Build, for example.
+        cmake_generator = os.environ.get("CMAKE_GENERATOR", "")
+
+        # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
+        cmake_args = [
+            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
+            f"-DPYTHON_EXECUTABLE={sys.executable}",
+            f"-DCMAKE_BUILD_TYPE={cfg}",  # not used on MSVC, but no harm
+            "-DSUPPORT_REGEX_LOOKAHEAD=ON",
+            "-DTOKENIZERS_BUILD_PYTHON=ON",
+            "-DCMAKE_POSITION_INDEPENDENT_CODE=ON",
+        ]
+        build_args = []
+
+        # Adding CMake arguments set as environment variable
+        # (needed e.g. to build for ARM OSX on conda-forge)
+        if "CMAKE_ARGS" in os.environ:
+            cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
+
+        if self.compiler.compiler_type != "msvc":
+            # Using Ninja-build since it a) is available as a wheel and b)
+            # multithreads automatically. MSVC would require all variables be
+            # exported for Ninja to pick it up, which is a little tricky to do.
+            # Users can override the generator with CMAKE_GENERATOR in CMake
+            # 3.15+.
+            if not cmake_generator or cmake_generator == "Ninja":
+                try:
+                    import ninja  # noqa: F401
+
+                    ninja_executable_path = os.path.join(ninja.BIN_DIR, "ninja")
+                    cmake_args += [
+                        "-GNinja",
+                        f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}",
+                    ]
+                except ImportError:
+                    pass
+
+        else:
+            # Single config generators are handled "normally"
+            single_config = any(x in cmake_generator for x in {"NMake", "Ninja"})
+
+            # CMake allows an arch-in-generator style for backward compatibility
+            contains_arch = any(x in cmake_generator for x in {"ARM", "Win64"})
+
+            # Specify the arch if using MSVC generator, but only if it doesn't
+            # contain a backward-compatibility arch spec already in the
+            # generator name.
+            if not single_config and not contains_arch:
+                cmake_args += ["-A", "x64"]
+
+            # Multi-config generators have a different way to specify configs
+            if not single_config:
+                cmake_args += [
+                    f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{cfg.upper()}={extdir}"
+                ]
+                build_args += ["--config", cfg]
+
+        if sys.platform.startswith("darwin"):
+            # Cross-compile support for macOS - respect ARCHFLAGS if set
+            archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
+            if archs:
+                cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
+
+        # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
+        # across all generators.
+        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
+            # self.parallel is a Python 3 only way to set parallel jobs by hand
+            # using -j in the build_ext call, not supported by pip or PyPA-build.
+            if hasattr(self, "parallel") and self.parallel:
+                # CMake 3.12+ only.
+                build_args += [f"-j{self.parallel}"]
+
+        build_temp = Path(self.build_temp) / ext.name
+        if not build_temp.exists():
+            build_temp.mkdir(parents=True)
+
+        subprocess.run(
+            ["cmake", ext.sourcedir] + cmake_args, cwd=build_temp, check=True
+        )
+        subprocess.run(
+            ["cmake", "--build", "."] + build_args, cwd=build_temp, check=True
+        )
+
 setup(
+    name="pytorch-tokenizers",
     version="0.1.0",
     long_description=long_description,
     long_description_content_type="text/markdown",
+    url="https://github.com/pytorch-labs/tokenizers",
     packages=find_packages(),
+    ext_modules=[CMakeExtension("pytorch_tokenizers_cpp")],
+    cmdclass={"build_ext": CMakeBuild},
+    zip_safe=False,
+    python_requires=">=3.8",
+    install_requires=[
+        "pybind11>=2.6.0",
+    ],
+    setup_requires=[
+        "pybind11>=2.6.0",
+        "cmake>=3.18",
+    ],
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: BSD License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Programming Language :: C++",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
 )