estimate.rs

crates/cortex-core/src/estimate.rs

1. //! Token estimates, for when the server reports none: most local servers, a
2. //! cancelled stream, an error part way. Also the composer's pre-send figure.
3. //! Anything built on these is marked "est." in the interface.
5. /// Roughly 4 characters per token for Latin script, 1 per CJK character.
6. pub fn estimate_tokens(text: &str) -> i64 {
7. let mut cjk = 0i64;
8. let mut other = 0i64;
9. for ch in text.chars() {
10. let c = ch as u32;
11. let wide = (0x3040..=0x30ff).contains(&c) // kana
12. || (0x3400..=0x9fff).contains(&c) // CJK unified
13. || (0xac00..=0xd7af).contains(&c) // hangul
14. || (0xf900..=0xfaff).contains(&c);
15. if wide {
16. cjk += 1;
17. } else {
18. other += 1;
19. }
20. }
21. cjk + (other + 3) / 4
22. }
24. /// Per-message overhead of the chat template.
25. pub const MESSAGE_OVERHEAD: i64 = 4;
26. /// A single image at the size Cortex sends it; real counts vary by model.
27. pub const IMAGE_TOKENS: i64 = 765;
29. /// A prompt of `(text, image_count)` messages.
30. pub fn estimate_prompt<'a>(messages: impl IntoIterator<Item = (&'a str, usize)>) -> i64 {
31. messages.into_iter().map(|(t, images)| estimate_tokens(t) + MESSAGE_OVERHEAD + images as i64 * IMAGE_TOKENS).sum()
32. }
34. #[cfg(test)]
35. mod tests {
36. use super::*;
38. #[test]
39. fn latin_and_cjk() {
40. assert_eq!(estimate_tokens(""), 0);
41. assert_eq!(estimate_tokens("abcd"), 1);
42. assert_eq!(estimate_tokens("abcde"), 2);
43. assert_eq!(estimate_tokens("\u{4f60}\u{597d}"), 2);
44. assert_eq!(estimate_prompt([("abcd", 0), ("", 1)]), 1 + 4 + 4 + 765);
45. }
46. }