feat: implement token counting from byte slices and add tokens module

Author: Arichy

src/language/language_type.rs

@@ -108,7 +108,7 @@ impl LanguageType {
                         // first character in the column, so removing starting whitespace
                         // could cause a miscount.
                         let line = if is_fortran { line } else { line.trim() };
-                        let tokens = Self::count_tokens(&String::from_utf8_lossy(line));
+                        let tokens = crate::tokens::count_tokens_from_bytes(line);
                         if line.trim().is_empty() {
                             (1, 0, 0, tokens)
                         } else if is_literate
@@ -134,11 +134,6 @@ impl LanguageType {
         }
     }
 
-    fn count_tokens(text: &str) -> usize {
-        let bpe = tiktoken_rs::p50k_base().unwrap();
-        bpe.encode_with_special_tokens(text).len()
-    }
-
     #[inline]
     fn parse_lines(
         self,
@@ -218,7 +213,7 @@ impl LanguageType {
             }
         }
 
-        let tokens = Self::count_tokens(&String::from_utf8_lossy(lines));
+        let tokens = crate::tokens::count_tokens_from_bytes(lines);
         stats.tokens += tokens;
 
         stats

src/lib.rs

@@ -54,6 +54,7 @@ mod consts;
 mod language;
 mod sort;
 mod stats;
+mod tokens;
 
 pub use self::{
     config::Config,

src/tokens.rs

@@ -0,0 +1,18 @@
+use once_cell::sync::Lazy;
+use tiktoken_rs::CoreBPE;
+
+static TOKENIZER: Lazy<CoreBPE> = Lazy::new(|| tiktoken_rs::p50k_base().unwrap());
+
+pub fn count_tokens(text: &str) -> usize {
+    TOKENIZER.encode_with_special_tokens(text).len()
+}
+
+pub fn count_tokens_from_bytes(bytes: &[u8]) -> usize {
+    match std::str::from_utf8(bytes) {
+        Ok(text) => count_tokens(text),
+        Err(_) => {
+            let text = String::from_utf8_lossy(bytes);
+            count_tokens(&text)
+        }
+    }
+}