add rust error handling

launcher/src/main.rs

@@ -139,6 +139,12 @@ struct Args {
     #[clap(long, env)]
     sharded: Option<bool>,
 
+    /// Whether this model is mean for embeddings or text generation.
+    /// By default models are for text generation.
+    /// Setting it to `true` will enable the embedding endpoints and disable the generation ones.
+    #[clap(long, env)]
+    embedding_model: Option<bool>,
+
     /// The number of shards to use if you don't want to use all GPUs on a given machine.
     /// You can use `CUDA_VISIBLE_DEVICES=0,1 lorax-launcher... --num_shard 2`
     /// and `CUDA_VISIBLE_DEVICES=2,3 lorax-launcher... --num_shard 2` to
@@ -1097,6 +1103,10 @@ fn spawn_webserver(
         router_args.push(origin.to_string());
     }
 
+    if args.embedding_model.unwrap_or(false) {
+        router_args.push("--embedding-model".to_string());
+    }
+
     // Ngrok
     if args.ngrok {
         router_args.push("--ngrok".to_string());

router/src/lib.rs

@@ -68,6 +68,8 @@ pub struct Info {
     pub docker_label: Option<&'static str>,
     #[schema(nullable = true, example = "http://localhost:8899")]
     pub request_logger_url: Option<String>,
+    #[schema(example = false)]
+    pub embedding_model: bool,
 }
 
 #[derive(Clone, Debug, Deserialize, ToSchema, Default)]

router/src/main.rs

@@ -64,6 +64,8 @@ struct Args {
     #[clap(long, env)]
     json_output: bool,
     #[clap(long, env)]
+    embedding_model: bool,
+    #[clap(long, env)]
     otlp_endpoint: Option<String>,
     #[clap(long, env)]
     cors_allow_origin: Option<Vec<String>>,
@@ -109,6 +111,7 @@ async fn main() -> Result<(), RouterError> {
         revision,
         validation_workers,
         json_output,
+        embedding_model,
         otlp_endpoint,
         cors_allow_origin,
         cors_allow_method,
@@ -372,6 +375,7 @@ async fn main() -> Result<(), RouterError> {
         ngrok_authtoken,
         ngrok_edge,
         adapter_source,
+        embedding_model,
     )
     .await?;
     Ok(())