Repository: keras-team/keras-nlp
Pull request #75 (2022-03-29 01:13:35)

Add tokenizer helper to convert tokens to ids
And from ids to tokens.

mattdangerw: @chenmoneygithub small change that I think will be useful for the masking layer, let me know what you think re-naming for these methods.

Patch:
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
new file mode 100644
index 00000000..48804717
--- /dev/null
+++ b/.github/workflows/publish-to-pypi.yml
@@ -0,0 +1,28 @@
+name: Publish to PyPI
+
+on: push
+jobs:
+  build-and-publish:
+    name: Build and publish to PyPI
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: Install pypa/build
+      run: >-
+        python -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: >-
+        python -m
+        build
+        --sdist
+        --wheel
+        --outdir dist/
+        .
+    - name: Publish distribution to PyPI
+      if: startsWith(github.ref, 'refs/tags')
+      uses: pypa/gh-action-pypi-publish@master
+      with:
+        password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.gitignore b/.gitignore
index 8aafcf6f..8a507ed0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,6 +8,7 @@ __pycache__/
 *.swo
 
 keras_nlp.egg-info/
+dist/
 
 .coverage
 .coverage.*
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 7e89ed49..fab19317 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -225,6 +225,17 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         """Get the size of the tokenizer vocabulary."""
         return len(self._vocab)
 
+    def id_to_token(self, id: int) -> str:
+        """Convert an integer id to a string token."""
+        return self._vocab[id]
+
+    def token_to_id(self, token: str) -> int:
+        """Convert a string token to an integer id."""
+        # This will be slow, but keep memory usage down compared to building a
+        # dict. Assuming the main use case is looking up a few special tokens
+        # early in the vocab, this should be fine.
+        return self._vocab.index(token)
+
     def get_config(self) -> Dict[str, Any]:
         config = super().get_config()
         config.update(
diff --git a/setup.py b/setup.py
index 177ce38f..bf41657d 100644
--- a/setup.py
+++ b/setup.py
@@ -14,14 +14,21 @@
 
 """Setup script."""
 
+import pathlib
+
 from setuptools import find_packages
 from setuptools import setup
 
+HERE = pathlib.Path(__file__).parent
+README = (HERE / "README.md").read_text()
+
 setup(
     name="keras-nlp",
     description=(
         "Industry-strength Natural Language Processing extensions for Keras."
     ),
+    long_description=README,
+    long_description_content_type="text/markdown",
     version="0.1.0",
     url="https://github.com/keras-team/keras-nlp",
     author="Keras team",


Test patch:
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index d57c5427..ab7fb1ea 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -66,6 +66,10 @@ class WordPieceTokenizerTest(tf.test.TestCase):
             tokenizer.get_vocabulary(),
             ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"],
         )
+        self.assertEqual(tokenizer.id_to_token(0), "[UNK]")
+        self.assertEqual(tokenizer.id_to_token(6), "fox")
+        self.assertEqual(tokenizer.token_to_id("[UNK]"), 0)
+        self.assertEqual(tokenizer.token_to_id("fox"), 6)
 
     def test_special_tokens(self):
         input_data = ["quick brown whale"]