refactor: streaming input text in dry-run mode (#154)

This commit is contained in:
sigoden 2023-10-27 09:33:34 +08:00 committed by GitHub
parent 2fad63be6a
commit 47818159a7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 15 additions and 2 deletions

View File

@ -17,6 +17,7 @@ use crate::{
client::localai::LocalAIClient,
config::{Config, SharedConfig},
repl::{ReplyStreamHandler, SharedAbortSignal},
utils::split_text,
};
#[allow(clippy::struct_excessive_bools)]
@ -92,7 +93,11 @@ pub trait Client {
tokio::select! {
ret = async {
if self.get_config().read().dry_run {
handler.text(&self.get_config().read().echo_messages(content))?;
let words = split_text(content)?;
for word in words {
tokio::time::sleep(Duration::from_millis(25)).await;
handler.text(&self.get_config().read().echo_messages(&word))?;
}
return Ok(());
}
self.send_message_streaming_inner(content, handler).await

View File

@ -412,7 +412,7 @@ impl Config {
if self.dry_run {
if let Ok(messages) = self.build_messages(input) {
let tokens = num_tokens_from_messages(&messages);
println!(">>> The following message consumes {tokens} tokens.");
println!(">>> This message consumes {tokens} tokens. <<<");
}
}
}

View File

@ -1,5 +1,6 @@
mod tiktoken;
use self::tiktoken::cl100k_base;
pub use self::tiktoken::{cl100k_base_singleton, count_tokens, text_to_tokens, tokens_to_text};
use arboard::Clipboard;
@ -41,3 +42,10 @@ pub fn copy(src: &str) -> Result<(), arboard::Error> {
let mut clipboard = Clipboard::new()?;
clipboard.set_text(src)
}
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
}