summaryrefslogtreecommitdiffstats
path: root/src/client/vertexai.rs
diff options
context:
space:
mode:
authorsigoden <sigoden@gmail.com>2024-03-06 08:35:40 +0800
committerGitHub <noreply@github.com>2024-03-06 08:35:40 +0800
commit8e5d4e55b1a5158a1f35adaf044dec545159f426 (patch)
treeb360cc4e9ad8acccea3c68b37057798cd1d657c0 /src/client/vertexai.rs
parentbe4e5e569a61c54d8ac8fb77144e9b1d01e3b81f (diff)
downloadaichat-8e5d4e55b1a5158a1f35adaf044dec545159f426.tar.gz
refactor: rename model's `max_tokens` to `max_input_tokens` (#339)
BREAKING CHANGE: rename model's `max_tokens` to `max_input_tokens`
Diffstat (limited to 'src/client/vertexai.rs')
-rw-r--r--src/client/vertexai.rs9
1 files changed, 5 insertions, 4 deletions
diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs
index 94becdb..0f92b4b 100644
--- a/src/client/vertexai.rs
+++ b/src/client/vertexai.rs
@@ -14,9 +14,10 @@ use serde::Deserialize;
use serde_json::{json, Value};
use std::path::PathBuf;
+// https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models
const MODELS: [(&str, usize, &str); 5] = [
- ("gemini-1.0-pro", 32760, "text"),
- ("gemini.1.0-pro-vision", 16384, "text,vision"),
+ ("gemini-1.0-pro", 24568, "text"),
+ ("gemini.1.0-pro-vision", 14336, "text,vision"),
("gemini-1.0-ultra", 8192, "text"),
("gemini.1.0-ultra-vision", 8192, "text,vision"),
("gemini-1.5-pro", 1000000, "text"),
@@ -66,10 +67,10 @@ impl VertexAIClient {
let client_name = Self::name(local_config);
MODELS
.into_iter()
- .map(|(name, max_tokens, capabilities)| {
+ .map(|(name, max_input_tokens, capabilities)| {
Model::new(client_name, name)
.set_capabilities(capabilities.into())
- .set_max_tokens(Some(max_tokens))
+ .set_max_input_tokens(Some(max_input_tokens))
.set_tokens_count_factors(TOKENS_COUNT_FACTORS)
})
.collect()