Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add ability to tokenize a string and return the decoded tokens using the correct BPE model #17

Merged
merged 3 commits into from Apr 16, 2023
Merged

Add ability to tokenize a string and return the decoded tokens using the correct BPE model #17

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c78268c
Add ability to tokenize a string and return the decoded tokens using …
jackbackes Apr 15, 2023
7ecf0a1
Add ChatCompletionRequestMessage Eq implementation
jackbackes Apr 16, 2023
d3461d3
Add split_by_token_with_special_tokens() method to CoreBPE
jackbackes Apr 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
26 changes: 26 additions & 0 deletions tiktoken-rs/src/vendor_tiktoken.rs
View file
Open in desktop
Expand Up @@ -235,6 +235,18 @@ impl CoreBPE {
ret
}

fn _decode_native_and_split(&self, tokens: &[usize]) -> Vec<Vec<u8>> {
let mut ret = Vec::with_capacity(tokens.len());
for token in tokens {
let token_bytes = self
.decoder
.get(token)
.unwrap_or_else(|| &self.special_tokens_decoder[token]);
ret.push(token_bytes.clone());
}
ret
}

fn _encode_ordinary_native(&self, text: &str) -> Vec<usize> {
// This is the core of the encoding logic; the other functions in here
// just make things complicated :-)
Expand Down Expand Up @@ -541,6 +553,20 @@ impl CoreBPE {
Err(e) => Err(anyhow!("Unable to decode into a valid UTF-8 string: {}", e)),
}
}

// tokenize a string and return the decoded tokens using the correct BPE model
// for example: "Hello world" -> ["Hello", " world"]
zurawiki marked this conversation as resolved.
Show resolved Hide resolved
pub fn split_by_token_with_special_tokens(&self, text: &str) -> Result<Vec<String>> {
// first, encode the text using the BPE model
let encoded = self.encode_with_special_tokens(text);

let tokenized = self._decode_native_and_split(&encoded);

tokenized
.iter()
.map(|token| String::from_utf8(token.clone()).map_err(|e| anyhow!(e.to_string())))
.collect()
}
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor nit: With the subfunctions you defined, I see some unnecessary clones and Vec collections and memory allocations. Consider inlining _decode_native_and_split or changing the function signatures to avoid clones

Suggested change
pub fn split_by_token_with_special_tokens(&self, text: &str) -> Result<Vec<String>> {
// first, encode the text using the BPE model
let encoded = self.encode_with_special_tokens(text);
let tokenized = self._decode_native_and_split(&encoded);
tokenized
.iter()
.map(|token| String::from_utf8(token.clone()).map_err(|e| anyhow!(e.to_string())))
.collect()
}
pub fn split_by_token_with_special_tokens(&self, text: &str) -> Result<Vec<String>> {
// first, encode the text using the BPE model
let encoded = self.encode_with_special_tokens(text);
encoded
.iter()
.map(|token| {
let token = self
.decoder
.get(token)
.unwrap_or_else(|| &self.special_tokens_decoder[token]);
String::from_utf8(token.clone()).map_err(|e| anyhow!(e.to_string()))
})
.collect()
}

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Disclaimer: I haven't thoroughly benchmarked the code so I'm not sure how big of an impact this would have on perf

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I may have gone too far in the other direction, but now the function returns an iterator and no work is done until the user runs it. Also eliminated unnecessary clones. It makes the api a little less straightforward but I don't think it's too bad.

}

#[cfg(feature = "python")]
Expand Down
7 changes: 7 additions & 0 deletions tiktoken-rs/tests/tiktoken.rs
View file
Open in desktop
Expand Up @@ -82,6 +82,13 @@ fn cl100k_base_test() {
);
}

#[test]
fn cl100k_split_test() {
let bpe = cl100k_base().unwrap();
let tokenized = bpe.split_by_token_with_special_tokens("This is a test with a lot of spaces").unwrap();
assert_eq!(tokenized, vec!["This", " is", " a", " test", " ", " with", " a", " lot", " of", " spaces"]);
}

#[test]
fn p50k_base_singleton_test() {
// let now = std::time::Instant::now();
Expand Down