diff options
Diffstat (limited to 'src/client/ernie.rs')
| -rw-r--r-- | src/client/ernie.rs | 26 |
1 files changed, 16 insertions, 10 deletions
diff --git a/src/client/ernie.rs b/src/client/ernie.rs index 0b10022..b0a0087 100644 --- a/src/client/ernie.rs +++ b/src/client/ernie.rs @@ -1,6 +1,6 @@ use super::{ - patch_system_message, Client, ErnieClient, ExtraConfig, Model, PromptType, ReplyHandler, - SendData, + patch_system_message, Client, ErnieClient, ExtraConfig, Model, ModelConfig, PromptType, + ReplyHandler, SendData, }; use crate::utils::PromptKind; @@ -73,6 +73,8 @@ pub struct ErnieConfig { pub name: Option<String>, pub api_key: Option<String>, pub secret_key: Option<String>, + #[serde(default)] + pub models: Vec<ModelConfig>, pub extra: Option<ExtraConfig>, } @@ -106,14 +108,18 @@ impl ErnieClient { pub fn list_models(local_config: &ErnieConfig) -> Vec<Model> { let client_name = Self::name(local_config); - MODELS - .into_iter() - .map(|(name, _, max_input_tokens, max_output_tokens)| { - Model::new(client_name, name) - .set_max_input_tokens(Some(max_input_tokens)) - .set_max_output_tokens(Some(max_output_tokens)) - }) // ERNIE tokenizer is different from cl100k_base - .collect() + if local_config.models.is_empty() { + MODELS + .into_iter() + .map(|(name, _, max_input_tokens, max_output_tokens)| { + Model::new(client_name, name) + .set_max_input_tokens(Some(max_input_tokens)) + .set_max_output_tokens(Some(max_output_tokens)) + }) // ERNIE tokenizer is different from cl100k_base + .collect() + } else { + Model::from_config(client_name, &local_config.models) + } } fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> { |