initial commit

Signed-off-by: Stephen Waits <steve@waits.net>

.build.yml

@@ -0,0 +1,73 @@
+image: archlinux
+packages:
+  - rustup
+  - cargo-audit
+  - cargo-deny
+  - cargo-msrv
+  - cargo-nextest
+  - cargo-outdated
+  - cargo-tarpaulin
+  - cargo-udeps
+secrets:
+  - 88cce185-d2ec-4573-9546-e33f36f79bbf
+sources:
+  - https://git.sr.ht/~swaits/bpe-tokenizer
+artifacts:
+  - bpe-tokenizer/tarpaulin-report.html
+tasks:
+  - setup: |
+      rustup toolchain install nightly stable
+      cd bpe-tokenizer/
+      rustup run stable cargo fetch
+  - format: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo fmt --verbose --check --all
+  - lint: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo clippy --verbose --all-targets --all-features
+  - outdated: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo outdated
+  - deny: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo deny check
+  - udeps: |
+      rustup default nightly
+      cd bpe-tokenizer/
+      cargo +nightly udeps --all-features
+  - audit: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo audit -d /tmp/advisory-db
+  - msrv: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo msrv verify
+  - coverage: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo tarpaulin --verbose --all-features --out html
+  - stable: |
+      rustup default stable
+      cd bpe-tokenizer/
+      cargo build --verbose --all-targets --all-features
+      cargo nextest run --verbose --all-targets --all-features
+  - nightly: |
+      rustup default nightly
+      cd bpe-tokenizer/
+      cargo clippy --verbose --all-targets --all-features  ||:
+      cargo build --verbose --all-targets --all-features   ||:
+      cargo nextest --verbose --all-targets --all-features ||:
+  - docs: |
+      cd bpe-tokenizer/
+      rustup run stable cargo doc --no-deps
+      rustup run nightly cargo doc --no-deps ||:
+  - mirror-to-github: |
+      cd ~/bpe-tokenizer
+      ssh-keyscan -t rsa github.com >> ~/.ssh/known_hosts
+      git remote add github git@github.com:swaits/bpe-tokenizer.git
+      git push --mirror github

.gitignore

@@ -0,0 +1,2 @@
+/target
+/Cargo.lock

Cargo.toml

@@ -0,0 +1,51 @@
+[package]
+name = "bpe-tokenizer"
+version = "0.1.0"
+authors = ["Stephen Waits <steve@waits.net>"]
+edition = "2021"
+description = "A BPE Tokenizer library."
+keywords = ["byte", "pair", "encoding", "bpe", "tokenizer"]
+categories = ["algorithms", "text-processing"]
+documentation = "https://docs.rs/bpe-tokenizer/"
+readme = "README.md"
+homepage = "https://github.com/swaits/bpe-tokenizer/"
+repository = "https://github.com/swaits/bpe-tokenizer/"
+rust-version = "1.64"
+license-file = "LICENSE"
+
+[dependencies]
+thiserror = "1.0.64"
+unicode-segmentation = "1.12.0"
+
+# optional, only needed when defaul-{small,medium,large} are enabled
+bincode = { version = "1.3.3", optional = true }
+lz4_flex = { version = "0.11.3", optional = true }
+serde = { version = "1.0.210", optional = true }
+
+[features]
+default-small = ["bincode", "lz4_flex", "serde"]
+default-medium = ["bincode", "lz4_flex", "serde"]
+default-large = ["bincode", "lz4_flex", "serde"]
+
+[build-dependencies]
+bincode = { version = "1.3.3", optional = true }
+lz4_flex = { version = "0.11.3", optional = true }
+serde = { version = "1.0.210", optional = true }
+
+[package.metadata.docs.rs]
+default = []
+all-features = true
+
+[profile.dev]
+opt-level = 1
+debug = true
+overflow-checks = true
+
+[profile.release]
+opt-level = 3
+lto = "thin"
+codegen-units = 1
+debug = false
+panic = "abort"
+overflow-checks = false
+strip = "symbols"

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Stephen Waits <steve@waits.net>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,143 @@
+# bpe-tokenizer
+
+A Rust implementation of Byte Pair Encoding (BPE) tokenization. This crate
+provides functionality to tokenize text into subword units using pre-trained
+vocabularies. BPE is widely used in natural language processing (NLP) tasks,
+where it breaks down words into subword tokens using a vocabulary of the most
+frequent token pairs.
+
+It supports **Unicode-aware** text segmentation for sentence and word splitting,
+making it suitable for processing a variety of languages and scripts.
+
+## Features
+
+- **Bring your own BPE token vocabularies**, or use ...
+- **Pre-trained multilingual vocabularies** sourced from the [BPEmb](https://github.com/bheinzerling/bpemb) project, with support for tokenizing text in **275 languages**.
+- **Unicode-aware sentence and word segmentation**: Leveraging the [`unicode-segmentation`](https://docs.rs/unicode-segmentation) crate for proper text splitting.
+
+## Installation
+
+To add this crate to your project, run:
+
+```bash
+cargo add bpe-tokenizer
+```
+
+Or manually include it in your `Cargo.toml`:
+
+```toml
+[dependencies]
+bpe-tokenizer = "<version>"
+```
+
+## Full Example
+
+Here is an example of how to create a `BytePairEncoder` from a string and use it
+to tokenize text:
+
+```rust
+use bpe_tokenizer::{BytePairEncoder, BytePairEncoderError};
+
+let vocab = BytePairEncoder::new_from_str("hello\t1\nworld\t2").unwrap();
+let tokenized = vocab.tokenize("Hello, world!");
+println!("{:?}", tokenized);
+```
+
+The output will be a vector of tokens:
+
+```text
+["<s>", "▁hello", "▁world", "</s>"]
+```
+
+Or load a vocabulary from a file:
+
+```rust
+use bpe_tokenizer::{BytePairEncoder, BytePairEncoderError};
+let vocab = BytePairEncoder::new_from_file("path/to/file.vocab").unwrap();
+```
+
+## Cargo Features
+
+The crate also includes several sizes of default pre-trained vocabularies, which
+are **optional** and can be enabled via Cargo features. They are sourced from
+Wikipedia data, pre-trained as part of the
+[BPEmb](https://github.com/bheinzerling/bpemb) project. These MIT-licensed
+vocabularies support 275 languages and provide different sizes depending on
+usage needs:
+
+### Available Optional Features
+
+- **`default-small` (100,000 tokens)**: Suitable for memory-constrained environments.
+- **`default-medium` (320,000 tokens)**: Balances between token coverage and memory efficiency.
+- **`default-large` (1,000,000 tokens)**: Provides the most detailed token representations for high granularity tasks.
+
+### Enabling Optional Features
+
+To use these default vocabularies, specify the feature in your `Cargo.toml`:
+
+```toml
+[dependencies]
+bpe-tokenizer = { version = "<version>", features = ["default-medium"] }
+```
+
+### Example with `default-medium` Vocabulary
+
+An example of using the **medium** vocabulary (320,000 tokens):
+
+```rust
+# #[cfg(feature = "default-medium")] {
+use bpe_tokenizer::{BytePairEncoder, BytePairEncoderError};
+
+let encoder = BytePairEncoder::new_default_medium().unwrap();
+let tokenized = encoder.tokenize("This is a test sentence.");
+println!("{:?}", tokenized);
+// Output: ["<s>", "▁this", "▁is", "▁a", "▁test", "▁sentence", "</s>"]
+# }
+```
+
+## Tokenization Functions
+
+The crate provides various ways to interact with the tokenizer:
+
+- **Tokenize into a flat `Vec<String>`**:
+
+  - `BytePairEncoder::tokenize`
+
+  Splits and flattens the text into tokens.
+
+  ```rust
+  let tokenized = vocab.tokenize("Example sentence.");
+  // Output: ["<s>", "▁example", "▁sentence", "</s>"]
+  ```
+
+- **Tokenize into nested sentence vectors `Vec<Vec<String>>`**:
+
+  - `BytePairEncoder::tokenize_sentences`
+
+  Useful for processing multiple sentences separately.
+
+  ```rust
+  let tokenized = vocab.tokenize_sentences("This is sentence one. And this is sentence two.");
+  // Output: [["<s>", "▁this", "▁is", "▁sentence", "▁one", "</s>"], ["<s>", "▁and", "▁this", "▁is", "▁sentence", "▁two", "</s>"]]
+  ```
+
+- **Iterative tokenization**:
+
+  - `BytePairEncoder::tokenize_iter` and `BytePairEncoder::tokenize_sentences_iter`
+
+  Provides an iterator over generated tokens for better memory efficiency in
+  large-scale text.
+
+  ```rust
+  let tokens_iter: Vec<String> = vocab.tokenize_iter("Example sentence").collect();
+  // Output: ["<s>", "▁example", "▁sentence", "</s>"]
+  ```
+
+## Licensing
+
+This crate is licensed under the [MIT License](LICENSE).
+
+## Contributing
+
+Contributions are welcome! Please open an issue, submit a pull request, or reach
+out if you'd like to contribute awesome new features or fixes to this crate.

build.rs

@@ -0,0 +1,75 @@
+#[cfg(any(
+    feature = "default-small",
+    feature = "default-medium",
+    feature = "default-large"
+))]
+use lz4_flex::block::compress_prepend_size;
+
+#[cfg(any(
+    feature = "default-small",
+    feature = "default-medium",
+    feature = "default-large"
+))]
+use std::{collections::HashMap, env, fs, path::PathBuf};
+
+fn main() {
+    #[cfg(feature = "default-small")]
+    process_vocab("multi.wiki.bpe.vs100000.vocab");
+
+    #[cfg(feature = "default-medium")]
+    process_vocab("multi.wiki.bpe.vs320000.vocab");
+
+    #[cfg(feature = "default-large")]
+    process_vocab("multi.wiki.bpe.vs1000000.vocab");
+}
+
+#[cfg(any(
+    feature = "default-small",
+    feature = "default-medium",
+    feature = "default-large"
+))]
+fn process_vocab(name: &str) {
+    // Path to the vocabulary file (ensure this path is correct)
+    let vocab_path = PathBuf::from(format!("vocab/{}", name));
+
+    // Load and parse the vocabulary into a HashMap.
+    let tokens: HashMap<String, isize> = load_vocab_hashmap(&vocab_path);
+
+    // Serialize the HashMap using bincode
+    let serialized = bincode::serialize(&tokens).unwrap();
+
+    // Compress the serialized data using zstd (with ultra compression level)
+    let compressed = compress_prepend_size(&serialized);
+
+    // Write the compressed data to a file in the build output directory
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+    let output_path = out_dir.join(format!("{}.hashmap.bincode.lz4", name));
+    fs::write(&output_path, compressed).unwrap();
+}
+
+#[cfg(any(
+    feature = "default-small",
+    feature = "default-medium",
+    feature = "default-large"
+))]
+fn load_vocab_hashmap(path: &PathBuf) -> HashMap<String, isize> {
+    // Read file's contents
+    let content = fs::read_to_string(path).unwrap();
+
+    let mut tokens = HashMap::new();
+
+    // Process each line in the file, each being a token-score pair
+    for line in content.lines() {
+        let (token, score_str) = match line.split_once('\t') {
+            Some(pair) => pair,
+            None => panic!("Invalid line in vocabulary file: {}", line),
+        };
+        let score = match score_str.parse::<isize>() {
+            Ok(score) => score,
+            Err(_) => panic!("Invalid score in vocabulary file: {}", line),
+        };
+        tokens.insert(token.to_string(), score);
+    }
+
+    tokens
+}