Unverified Commit 67bc6a1f authored by MatejKosec's avatar MatejKosec Committed by GitHub
Browse files

fix(discovery): add missing Audios model type to is_model_type_list_empty (#8241)


Signed-off-by: default avatarMatej Kosec <mkosec@nvidia.com>
parent 63787939
...@@ -218,6 +218,14 @@ impl ModelManager { ...@@ -218,6 +218,14 @@ impl ModelManager {
.collect() .collect()
} }
pub fn list_audios_models(&self) -> Vec<String> {
self.models
.iter()
.filter(|entry| entry.value().has_audios_engine())
.map(|entry| entry.key().clone())
.collect()
}
pub fn list_videos_models(&self) -> Vec<String> { pub fn list_videos_models(&self) -> Vec<String> {
self.models self.models
.iter() .iter()
......
...@@ -108,6 +108,8 @@ fn is_model_type_list_empty(manager: &ModelManager, model_type: ModelType) -> bo ...@@ -108,6 +108,8 @@ fn is_model_type_list_empty(manager: &ModelManager, model_type: ModelType) -> bo
manager.list_embeddings_models().is_empty() manager.list_embeddings_models().is_empty()
} else if model_type == ModelType::Images { } else if model_type == ModelType::Images {
manager.list_images_models().is_empty() manager.list_images_models().is_empty()
} else if model_type == ModelType::Audios {
manager.list_audios_models().is_empty()
} else if model_type == ModelType::Videos { } else if model_type == ModelType::Videos {
manager.list_videos_models().is_empty() manager.list_videos_models().is_empty()
} else if model_type == ModelType::TensorBased { } else if model_type == ModelType::TensorBased {
...@@ -1059,6 +1061,7 @@ mod tests { ...@@ -1059,6 +1061,7 @@ mod tests {
assert!(is_model_type_list_empty(&mm, ModelType::Completions)); assert!(is_model_type_list_empty(&mm, ModelType::Completions));
assert!(is_model_type_list_empty(&mm, ModelType::Embedding)); assert!(is_model_type_list_empty(&mm, ModelType::Embedding));
assert!(is_model_type_list_empty(&mm, ModelType::Images)); assert!(is_model_type_list_empty(&mm, ModelType::Images));
assert!(is_model_type_list_empty(&mm, ModelType::Audios));
assert!(is_model_type_list_empty(&mm, ModelType::Videos)); assert!(is_model_type_list_empty(&mm, ModelType::Videos));
assert!(is_model_type_list_empty(&mm, ModelType::TensorBased)); assert!(is_model_type_list_empty(&mm, ModelType::TensorBased));
assert!(is_model_type_list_empty(&mm, ModelType::Prefill)); assert!(is_model_type_list_empty(&mm, ModelType::Prefill));
...@@ -1076,6 +1079,7 @@ mod tests { ...@@ -1076,6 +1079,7 @@ mod tests {
assert!(is_model_type_list_empty(&mm, ModelType::Completions)); assert!(is_model_type_list_empty(&mm, ModelType::Completions));
assert!(is_model_type_list_empty(&mm, ModelType::Embedding)); assert!(is_model_type_list_empty(&mm, ModelType::Embedding));
assert!(is_model_type_list_empty(&mm, ModelType::Images)); assert!(is_model_type_list_empty(&mm, ModelType::Images));
assert!(is_model_type_list_empty(&mm, ModelType::Audios));
assert!(is_model_type_list_empty(&mm, ModelType::Videos)); assert!(is_model_type_list_empty(&mm, ModelType::Videos));
assert!(is_model_type_list_empty(&mm, ModelType::TensorBased)); assert!(is_model_type_list_empty(&mm, ModelType::TensorBased));
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment