From f6da06dad9b2a76016209a7d58ad923c2c72f150 Mon Sep 17 00:00:00 2001 From: sigoden Date: Wed, 1 Nov 2023 22:15:55 +0800 Subject: refactor: improve code quanity (#194) - extends ModelInfo for tokens calculating - refactor config/session.rs, improve export, render, getter/setter - modify main.rs, allow --model override session.model --- src/config/message.rs | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) (limited to 'src/config/message.rs') diff --git a/src/config/message.rs b/src/config/message.rs index 5882337..55b2663 100644 --- a/src/config/message.rs +++ b/src/config/message.rs @@ -1,5 +1,3 @@ -use crate::utils::count_tokens; - use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Deserialize, Serialize)] @@ -25,22 +23,19 @@ pub enum MessageRole { User, } +#[allow(dead_code)] impl MessageRole { - #[allow(dead_code)] pub fn is_system(&self) -> bool { matches!(self, MessageRole::System) } -} -pub fn num_tokens_from_messages(messages: &[Message]) -> usize { - let mut num_tokens = 0; - for message in messages.iter() { - num_tokens += 4; - num_tokens += count_tokens(&message.content); - num_tokens += 1; // role always take 1 token + pub fn is_user(&self) -> bool { + matches!(self, MessageRole::User) + } + + pub fn is_assistant(&self) -> bool { + matches!(self, MessageRole::Assistant) } - num_tokens += 2; - num_tokens } #[cfg(test)] -- cgit v1.2.3