feat: improve Tokenizers (#45) (#46)

* feat: improve tokenizers

* feat: id to token & token to id

crates/tokenizers/Cargo.toml

@@ -11,6 +11,7 @@ ndarray = "0.15.6"
 ndarray-rand = "0.14.0"
 serde = {version = "1.0", features = ["derive"]}
 serde_json = "1.0"
+serde-wasm-bindgen = "0.6.0"
 tokenizers = { version="0.14.1", default-features=false, features = ["unstable_wasm"]}
 wasm-bindgen = "=0.2.87"
 getrandom = { version = "0.2", features = ["js"] }

crates/tokenizers/src/wasm.rs

@@ -1,5 +1,5 @@
 use crate::RESOURCES;
-use std::str::FromStr;
+use std::{collections::HashMap, str::FromStr};
 use tokenizers::{models::bpe::BPE, tokenizer::Tokenizer};
 use wasm_bindgen::prelude::*;
 
@@ -39,7 +39,7 @@ pub fn wasm_bpe_default() -> usize {
 }
 
 #[wasm_bindgen]
-pub fn wasm_tokenizer_tokenize(id: usize, string: String) -> Vec<u32> {
+pub fn wasm_tokenizer_encode(id: usize, string: String) -> Vec<u32> {
     let mut data: Vec<u32> = Vec::new();
     RESOURCES.with(|cell| {
         let tokenizers = cell.tokenizer.borrow_mut();
@@ -53,3 +53,53 @@ pub fn wasm_tokenizer_tokenize(id: usize, string: String) -> Vec<u32> {
     });
     data
 }
+
+#[wasm_bindgen]
+pub fn wasm_tokenizer_get_vocab(id: usize, with_added_tokens: bool) -> JsValue {
+    let mut data: HashMap<String, u32> = HashMap::new();
+    RESOURCES.with(|cell| {
+        let tokenizers = cell.tokenizer.borrow_mut();
+        data = tokenizers[id].get_vocab(with_added_tokens)
+    });
+    serde_wasm_bindgen::to_value(&data).unwrap()
+}
+
+#[wasm_bindgen]
+pub fn wasm_tokenizer_get_vocab_size(id: usize, with_added_tokens: bool) -> usize {
+    let mut data: usize = 0;
+    RESOURCES.with(|cell| {
+        let tokenizers = cell.tokenizer.borrow_mut();
+        data = tokenizers[id].get_vocab_size(with_added_tokens)
+    });
+    data
+}
+
+#[wasm_bindgen]
+pub fn wasm_tokenizer_decode(id: usize, ids: &[u32], skip_special_tokens: bool) -> String {
+    let mut data: String = String::new();
+    RESOURCES.with(|cell| {
+        let tokenizers = cell.tokenizer.borrow_mut();
+        data = tokenizers[id].decode(ids, skip_special_tokens).unwrap()
+    });
+    data
+}
+
+#[wasm_bindgen]
+pub fn wasm_tokenizer_token_to_id(id: usize, token: String) -> u32 {
+    let mut data: u32 = 0;
+    RESOURCES.with(|cell| {
+        let tokenizers = cell.tokenizer.borrow_mut();
+        data = tokenizers[id].token_to_id(token.as_str()).unwrap()
+    });
+    data
+}
+
+#[wasm_bindgen]
+pub fn wasm_tokenizer_id_to_token(id: usize, token_id: u32) -> String {
+    let mut data: String = String::new();
+    RESOURCES.with(|cell| {
+        let tokenizers = cell.tokenizer.borrow_mut();
+        data = tokenizers[id].id_to_token(token_id).unwrap()
+    });
+    data
+}

examples/tokenizers/basic.ts

@@ -2,10 +2,13 @@ import { init, Tokenizer } from "../../tokenizers/mod.ts";
 
 await init();
 
-const tokenizer = Tokenizer.fromJson(
+const tokenizer = Tokenizer.fromJSON(
   await (await fetch(
     `https://huggingface.co/satvikag/chatbot/resolve/main/tokenizer.json`,
   )).text(),
 );
 
-console.log(tokenizer.tokenize("Hello World!"));
+const encoded = tokenizer.encode("Hello World!");
+console.log(encoded);
+const decoded = tokenizer.decode(encoded);
+console.log(decoded);