From 5c0383f908eaee86539103b536a95bd94d32d401 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 30 Oct 2023 16:32:11 +0800
Subject: fix: dry run on role or session (#181)

---
 src/utils/mod.rs | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

(limited to 'src/utils/mod.rs')

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 699f8a0..db53fd6 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -1,10 +1,8 @@
 mod tiktoken;
 
-use self::tiktoken::cl100k_base;
-pub use self::tiktoken::{cl100k_base_singleton, count_tokens, text_to_tokens, tokens_to_text};
+pub use self::tiktoken::cl100k_base_singleton;
 
 use chrono::prelude::*;
-use crossterm::style::{Color, Stylize};
 use std::io::{stdout, Write};
 
 #[macro_export]
@@ -32,16 +30,18 @@ pub fn get_env_name(key: &str) -> String {
     )
 }
 
-#[allow(unused)]
-pub fn emphasis(text: &str) -> String {
-    text.stylize().with(Color::White).to_string()
+/// Split text to tokens
+pub fn tokenize(text: &str) -> Vec<String> {
+    let tokens = cl100k_base_singleton().lock().tokenize(text);
+    tokens.into_iter().map(|(_, text)| text).collect()
 }
 
-pub fn split_text(text: &str) -> Result<Vec<String>, anyhow::Error> {
-    let bpe = cl100k_base()?;
-    let tokens = bpe.encode_with_special_tokens(text);
-    let data: Result<Vec<String>, _> = tokens.into_iter().map(|v| bpe.decode(&[v])).collect();
-    data
+/// Count how many tokens a piece of text needs to consume
+pub fn count_tokens(text: &str) -> usize {
+    cl100k_base_singleton()
+        .lock()
+        .encode_with_special_tokens(text)
+        .len()
 }
 
 pub fn light_theme_from_colorfgbg(colorfgbg: &str) -> Option<bool> {
-- 
cgit v1.2.3