From 3a3388375be05758d5f5574cce1d0f8eb7bdd604 Mon Sep 17 00:00:00 2001 From: sigoden Date: Wed, 4 Dec 2024 21:03:59 +0800 Subject: refactor: improve retrieve model (#1036) - check the model type while retrieve model - select chat/reranker model even if it is missed in client models - find predefined-models for openai-compatible client with startsWith - remove client::ApiType --- src/serve.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) (limited to 'src/serve.rs') diff --git a/src/serve.rs b/src/serve.rs index c19be1f..b3fe5e0 100644 --- a/src/serve.rs +++ b/src/serve.rs @@ -75,7 +75,7 @@ impl Server { fn new(config: &GlobalConfig) -> Self { let mut config = config.read().clone(); config.functions = Functions::default(); - let mut models = list_models(&config); + let mut models = list_all_models(&config); let mut default_model = config.model.clone(); default_model.data_mut().name = DEFAULT_MODEL_NAME.into(); models.insert(0, &default_model); @@ -483,7 +483,8 @@ impl Server { let config = Arc::new(RwLock::new(self.config.clone())); - let embedding_model = Model::retrieve_embedding(&config.read(), &embedding_model_id)?; + let embedding_model = + Model::retrieve_model(&config.read(), &embedding_model_id, ModelType::Embedding)?; let texts = match input { EmbeddingsReqBodyInput::Single(v) => vec![v], @@ -542,7 +543,8 @@ impl Server { let config = Arc::new(RwLock::new(self.config.clone())); - let reranker_model = Model::retrieve_embedding(&config.read(), &reranker_model_id)?; + let reranker_model = + Model::retrieve_model(&config.read(), &reranker_model_id, ModelType::Reranker)?; let client = init_client(&config, Some(reranker_model))?; let data = client -- cgit v1.2.3