huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 04:14:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: prefer env value from clap for better defaults

Browse Source
This commit is contained in:
drbh 2024-01-18 11:03:05 -05:00
parent 90541fba07
commit 215afc15f0
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
router/src/main.rs
View File

@ -71,6 +71,8 @@ struct Args {
ngrok_authtoken: Option<String>, ngrok_authtoken: Option<String>,
#[clap(long, env)] #[clap(long, env)]
ngrok_edge: Option<String>, ngrok_edge: Option<String>,
#[clap(long, env, default_value_t = false)]
chat_enabled_api: bool,
} }
#[tokio::main] #[tokio::main]
@ -102,6 +104,7 @@ async fn main() -> Result<(), RouterError> {
ngrok, ngrok,
ngrok_authtoken, ngrok_authtoken,
ngrok_edge, ngrok_edge,
chat_enabled_api,
} = args; } = args;
// Launch Tokio runtime // Launch Tokio runtime
@ -345,6 +348,7 @@ async fn main() -> Result<(), RouterError> {
ngrok_authtoken, ngrok_authtoken,
ngrok_edge, ngrok_edge,
tokenizer_config, tokenizer_config,
chat_enabled_api,
) )
.await?; .await?;
Ok(()) Ok(())

3
router/src/server.rs
View File

@ -708,6 +708,7 @@ pub async fn run(
ngrok_authtoken: Option<String>, ngrok_authtoken: Option<String>,
ngrok_edge: Option<String>, ngrok_edge: Option<String>,
tokenizer_config: HubTokenizerConfig, tokenizer_config: HubTokenizerConfig,
chat_enabled_api: bool,
) -> Result<(), axum::BoxError> { ) -> Result<(), axum::BoxError> {
// OpenAPI documentation // OpenAPI documentation
#[derive(OpenApi)] #[derive(OpenApi)]
@ -871,7 +872,7 @@ pub async fn run(
.route("/metrics", get(metrics)); .route("/metrics", get(metrics));
// Conditional AWS Sagemaker route // Conditional AWS Sagemaker route
let aws_sagemaker_route = if std::env::var("OAI_ENABLED").map_or(false, |val| val == "true") { let aws_sagemaker_route = if chat_enabled_api {
Router::new().route("/invocations", post(chat_completions)) // Use 'chat_completions' for OAI_ENABLED Router::new().route("/invocations", post(chat_completions)) // Use 'chat_completions' for OAI_ENABLED
} else { } else {
Router::new().route("/invocations", post(compat_generate)) // Use 'compat_generate' otherwise Router::new().route("/invocations", post(compat_generate)) // Use 'compat_generate' otherwise