From 617f5000394ce27f8e536a67f057e581933d3a04 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 1 Aug 2024 21:01:23 -0400 Subject: [PATCH] fix model weight search --- .../searches/model_weights/search.json | 7 + .../production/elasticsearch-cli/src/main.rs | 17 +- .../src/searches/search_model_weights.rs | 419 +++++++++++++++--- .../weights/search_model_weights_handler.rs | 30 +- 4 files changed, 409 insertions(+), 64 deletions(-) diff --git a/_elasticsearch/searches/model_weights/search.json b/_elasticsearch/searches/model_weights/search.json index dada571316..5d1f92d3b8 100644 --- a/_elasticsearch/searches/model_weights/search.json +++ b/_elasticsearch/searches/model_weights/search.json @@ -4,6 +4,13 @@ "must": [ { "bool": { + "must" : [ + { + "term": { + "is_deleted": false + } + } + ], "should": [ { "fuzzy": { diff --git a/crates/cli/production/elasticsearch-cli/src/main.rs b/crates/cli/production/elasticsearch-cli/src/main.rs index 6d7565c95b..3a67c19b44 100644 --- a/crates/cli/production/elasticsearch-cli/src/main.rs +++ b/crates/cli/production/elasticsearch-cli/src/main.rs @@ -5,6 +5,9 @@ //! //! The intent is to be able to quickly populate documents and indices. //! + +use std::collections::HashSet; +use std::iter::FromIterator; use elasticsearch::Elasticsearch; use elasticsearch::http::transport::Transport; use log::info; @@ -12,7 +15,7 @@ use sqlx::{MySql, Pool}; use sqlx::mysql::MySqlPoolOptions; use config::shared_constants::DEFAULT_RUST_LOG; -use elasticsearch_schema::searches::search_model_weights::search_model_weights; +use elasticsearch_schema::searches::search_model_weights::{search_model_weights, SearchArgs}; use elasticsearch_schema::searches::search_tts_models::search_tts_models; use enums::by_table::model_weights::weights_types::WeightsType; use errors::AnyhowResult; @@ -58,8 +61,16 @@ pub async fn main() -> AnyhowResult<()> { } Action::SearchModelWeights => { info!("Searching model weights..."); - //let results = search_model_weights(&elasticsearch, "zel", Some("en"), Some(WeightsType::SoVitsSvc)).await?; - let results = search_model_weights(&elasticsearch, "zel", None, Some(WeightsType::RvcV2), None).await?; + + let results = search_model_weights(SearchArgs { + search_term: "zel", + maybe_creator_user_token: None, + maybe_ietf_primary_language_subtag: None, + maybe_weights_categories: None, + maybe_weights_types: Some(HashSet::from_iter(vec![WeightsType::RvcV2])), + client: &elasticsearch, + }).await?; + for result in results { println!("Result: {:#?}", result); } diff --git a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs index 32f227c1c3..b81a04c98a 100644 --- a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs +++ b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs @@ -1,12 +1,25 @@ +use std::collections::HashSet; + use elasticsearch::{Elasticsearch, SearchParts}; +use once_cell::sync::Lazy; use serde_json::{json, Value}; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use errors::{anyhow, AnyhowResult}; +use tokens::tokens::users::UserToken; use crate::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument}; +static JSON_QUERY : Lazy = Lazy::new(|| { + const QUERY_TEMPLATE : &str = include_str!("../../../../../_elasticsearch/searches/model_weights/search.json"); + + let json : Value = serde_json::from_str(QUERY_TEMPLATE) + .expect("json should parse"); + + json +}); + pub struct SearchModelWeightsQuery<'a> { pub search_term: &'a str, pub maybe_language_subtag: Option<&'a str>, @@ -103,26 +116,23 @@ impl <'a>SearchModelWeightsQuery<'a> { } } -pub async fn search_model_weights( - client: &Elasticsearch, - search_term: &str, - maybe_language_subtag: Option<&str>, - maybe_weights_type: Option, - maybe_weights_category: Option, -) -> AnyhowResult> { +pub struct SearchArgs<'a> { + pub search_term: &'a str, + // pub is_featured: Option, + pub maybe_creator_user_token: Option<&'a UserToken>, + pub maybe_ietf_primary_language_subtag: Option<&'a str>, + pub maybe_weights_categories: Option>, + pub maybe_weights_types: Option>, - let query = SearchModelWeightsQuery { - search_term, - maybe_language_subtag, - maybe_weights_type, - maybe_weights_category, - }; + pub client: &'a Elasticsearch, +} - let search_json = query.query()?; +pub async fn search_model_weights(args: SearchArgs<'_>) -> AnyhowResult> { + let query = build_query(&args)?; - let search_response = client + let search_response = args.client .search(SearchParts::Index(&[MODEL_WEIGHT_INDEX])) - .body(search_json) + .body(query) .size(30) .allow_no_indices(true) .send() @@ -159,41 +169,341 @@ pub async fn search_model_weights( Ok(documents) } +fn build_query(args: &SearchArgs) -> AnyhowResult { + let query = JSON_QUERY.clone(); + + let query = jsonpath_lib::replace_with(query, "$.query.bool.must[0].bool.must", &mut |_| { + let mut predicates = vec![ + must_be_not_deleted(), + ]; + +// if let Some(is_featured) = args.is_featured { +// predicates.push(featured_predicate(is_featured)); +// } +// + + if let Some(language_subtag) = args.maybe_ietf_primary_language_subtag { + predicates.push(language_subtag_predicate(language_subtag)); + } + + if let Some(creator_user_token) = args.maybe_creator_user_token { + predicates.push(creator_user_token_predicate(creator_user_token)); + } + + if let Some(weights_categories) = &args.maybe_weights_categories { + predicates.push(weights_categories_predicates(weights_categories)); + } + + if let Some(weights_types) = &args.maybe_weights_types { + predicates.push(weights_types_predicates(weights_types)); + } + + Some(json!(predicates)) + })?; + + let query = jsonpath_lib::replace_with(query, "$.query.bool.must[0].bool.should[0].fuzzy.title.value", &mut |_| { + Some(json!(args.search_term)) + })?; + + let query = jsonpath_lib::replace_with(query, "$.query.bool.must[0].bool.should[1].match.title.query", &mut |_| { + Some(json!(args.search_term)) + })?; + + let query = jsonpath_lib::replace_with(query, "$.query.bool.must[0].bool.should[2].multi_match.query", &mut |_| { + Some(json!(args.search_term)) + })?; + + Ok(query) +} + +fn must_be_not_deleted() -> Value { + json!({ + "term": { + "is_deleted": false, + } + }) +} + +// fn featured_predicate(is_featured: bool) -> Value { +// json!({ +// "term": { +// "is_featured": is_featured, +// } +// }) +// } + +fn creator_user_token_predicate(creator_user_token: &UserToken) -> Value { + json!({ + "term": { + "maybe_creator_user_token": creator_user_token.as_str(), + } + }) +} + +fn language_subtag_predicate(language_subtag: &str) -> Value { + json!({ + "term": { + "maybe_ietf_primary_language_subtag": language_subtag, + } + }) +} + + +fn weights_categories_predicates(weights_categories: &HashSet) -> Value { + should_predicates(weights_categories.iter() + .map(|weight_category| { + json!({ + "term": { + "weights_category": weight_category.to_str(), + } + }) + }) + .collect()) +} + +fn weights_types_predicates(weights_types: &HashSet) -> Value { + should_predicates(weights_types.iter() + .map(|weights_type| { + json!({ + "term": { + "weights_type": weights_type.to_str(), + } + }) + }) + .collect()) +} + +// NB: "Should" is a logical OR. +fn should_predicates(predicates: Vec) -> Value { + json!({ + "bool": { + "should": predicates, + } + }) +} + + #[cfg(test)] mod tests { - use regex::Regex; + use std::collections::HashSet; + use std::iter::FromIterator; + + use serde_json::Value; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; - use errors::AnyhowResult; + use tokens::tokens::users::UserToken; - use crate::searches::search_model_weights::SearchModelWeightsQuery; + use crate::searches::search_model_weights::build_query; + use crate::searches::search_model_weights::SearchArgs; - fn query_to_json_string(query: SearchModelWeightsQuery<'_>) -> AnyhowResult { - let json = query.query()?; - let json = serde_json::to_string(&json)?; - Ok(json) - } + #[test] + fn test_default_search() { + let search = build_query(&SearchArgs { + search_term: "foo", + maybe_creator_user_token: None, + maybe_ietf_primary_language_subtag: None, + maybe_weights_categories: None, + maybe_weights_types: None, + client: &elasticsearch::Elasticsearch::default(), + }).unwrap(); - fn compact_json(json: &str) -> String { - let regex = Regex::new("\\s+").expect("regex should parse"); - let json = regex.replace_all(&json, ""); - json.to_string() + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.must[0].term.is_deleted").unwrap(); + + assert_eq!(value[0], &Value::Bool(false)); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[0].fuzzy.title.value").unwrap(); + + assert_eq!(value[0], &Value::String("foo".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[1].match.title.query").unwrap(); + + assert_eq!(value[0], &Value::String("foo".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[2].multi_match.query").unwrap(); + + assert_eq!(value[0], &Value::String("foo".to_string())); } #[test] - fn default_search_only_keyword() { - let query = SearchModelWeightsQuery { - search_term: "FOO_BAR_BAZ", - maybe_language_subtag: None, - maybe_weights_type: None, - maybe_weights_category: None, - }; + fn test_creator_user_token() { + let search = build_query(&SearchArgs { + search_term: "asdf", + maybe_creator_user_token: Some(&UserToken::new_from_str("USER_TOKEN")), + maybe_ietf_primary_language_subtag: None, + maybe_weights_categories: None, + maybe_weights_types: None, + client: &elasticsearch::Elasticsearch::default(), + }).unwrap(); - let json = query_to_json_string(query).unwrap(); + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.must[0].term.is_deleted").unwrap(); - // NB: Keys in emitted JSON are sorted. - let expected_json = r#" + assert_eq!(value[0], &Value::Bool(false)); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.must[1].term.maybe_creator_user_token").unwrap(); + + assert_eq!(value[0], &Value::String("USER_TOKEN".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[0].fuzzy.title.value").unwrap(); + + assert_eq!(value[0], &Value::String("asdf".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[1].match.title.query").unwrap(); + + assert_eq!(value[0], &Value::String("asdf".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[2].multi_match.query").unwrap(); + + assert_eq!(value[0], &Value::String("asdf".to_string())); + } + + #[test] + fn test_ietf_primary_language_subtag() { + let search = build_query(&SearchArgs { + search_term: "foo", + maybe_creator_user_token: None, + maybe_ietf_primary_language_subtag: Some("ja"), + maybe_weights_categories: None, + maybe_weights_types: None, + client: &elasticsearch::Elasticsearch::default(), + }).unwrap(); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.must[0].term.is_deleted").unwrap(); + + assert_eq!(value[0], &Value::Bool(false)); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.must[1].term.maybe_ietf_primary_language_subtag").unwrap(); + + assert_eq!(value[0], &Value::String("ja".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[0].fuzzy.title.value").unwrap(); + + assert_eq!(value[0], &Value::String("foo".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[1].match.title.query").unwrap(); + + assert_eq!(value[0], &Value::String("foo".to_string())); + + let value = jsonpath_lib::select( + &search, "$.query.bool.must[0].bool.should[2].multi_match.query").unwrap(); + + assert_eq!(value[0], &Value::String("foo".to_string())); + } + + #[test] + fn test_weight_types() { + let search = build_query(&SearchArgs { + search_term: "bar", + maybe_creator_user_token: None, + maybe_ietf_primary_language_subtag: None, + maybe_weights_categories: None, + maybe_weights_types: Some(HashSet::from_iter(vec![ + WeightsType::Tacotron2, + WeightsType::GptSoVits, + ])), + client: &elasticsearch::Elasticsearch::default(), + }).unwrap(); + + let value = select(&search, "$.query.bool.must[0].bool.must[0].term.is_deleted"); + + assert_eq!(value[0], &Value::Bool(false)); + + let values = select_str_values( + &search, "$.query.bool.must[0].bool.must[1].bool.should[*].term.weights_type"); + + assert_eq!(values.len(), 2); + assert!(values.contains(&"tt2")); + assert!(values.contains(&"gpt_so_vits")); + } + + #[test] + fn test_weight_categories() { + let search = build_query(&SearchArgs { + search_term: "bar", + maybe_creator_user_token: None, + maybe_ietf_primary_language_subtag: None, + maybe_weights_types: None, + maybe_weights_categories: Some(HashSet::from_iter(vec![ + WeightsCategory::TextToSpeech, + WeightsCategory::VoiceConversion, + ])), + client: &elasticsearch::Elasticsearch::default(), + }).unwrap(); + + let value = select(&search, "$.query.bool.must[0].bool.must[0].term.is_deleted"); + + assert_eq!(value[0], &Value::Bool(false)); + + let values = select_str_values( + &search, "$.query.bool.must[0].bool.must[1].bool.should[*].term.weights_category"); + + assert_eq!(values.len(), 2); + assert!(values.contains(&"text_to_speech")); + assert!(values.contains(&"voice_conversion")); + } + + fn select<'a>(search: &'a Value, path: &str) -> Vec<&'a Value> { + jsonpath_lib::select(search, path).unwrap() + } + + fn select_str_values<'a>(search: &'a Value, path: &str) -> Vec<&'a str> { + select(search, path).into_iter() + .map(|value| { + match value { + Value::String(inner) => inner.as_str(), + _ => panic!("Expected string"), + } + }) + .collect() + } + + mod old_tests { + use regex::Regex; + + use enums::by_table::model_weights::weights_category::WeightsCategory; + use enums::by_table::model_weights::weights_types::WeightsType; + use errors::AnyhowResult; + + use crate::searches::search_model_weights::SearchModelWeightsQuery; + + fn query_to_json_string(query: SearchModelWeightsQuery<'_>) -> AnyhowResult { + let json = query.query()?; + let json = serde_json::to_string(&json)?; + Ok(json) + } + + fn compact_json(json: &str) -> String { + let regex = Regex::new("\\s+").expect("regex should parse"); + let json = regex.replace_all(&json, ""); + json.to_string() + } + + #[test] + fn default_search_only_keyword() { + let query = SearchModelWeightsQuery { + search_term: "FOO_BAR_BAZ", + maybe_language_subtag: None, + maybe_weights_type: None, + maybe_weights_category: None, + }; + + let json = query_to_json_string(query).unwrap(); + + // NB: Keys in emitted JSON are sorted. + let expected_json = r#" { "query": { "bool": { @@ -238,24 +548,24 @@ mod tests { } "#; - let expected_json = compact_json(&expected_json); + let expected_json = compact_json(&expected_json); - assert_eq!(&json, &expected_json); - } + assert_eq!(&json, &expected_json); + } - #[test] - fn search_with_language_and_weights() { - let query = SearchModelWeightsQuery { - search_term: "FOO_BAR_BAZ", - maybe_language_subtag: Some("en"), - maybe_weights_type: Some(WeightsType::SoVitsSvc), - maybe_weights_category: Some(WeightsCategory::Vocoder), - }; + #[test] + fn search_with_language_and_weights() { + let query = SearchModelWeightsQuery { + search_term: "FOO_BAR_BAZ", + maybe_language_subtag: Some("en"), + maybe_weights_type: Some(WeightsType::SoVitsSvc), + maybe_weights_category: Some(WeightsCategory::Vocoder), + }; - let json = query_to_json_string(query).unwrap(); + let json = query_to_json_string(query).unwrap(); - // NB: Keys in emitted JSON are sorted. - let expected_json = r#" + // NB: Keys in emitted JSON are sorted. + let expected_json = r#" { "query": { "bool": { @@ -315,8 +625,9 @@ mod tests { } "#; - let expected_json = compact_json(&expected_json); + let expected_json = compact_json(&expected_json); - assert_eq!(&json, &expected_json); + assert_eq!(&json, &expected_json); + } } } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs index 407cee75dc..cf4c5db3af 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs @@ -15,7 +15,7 @@ use log::error; use utoipa::ToSchema; use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; -use elasticsearch_schema::searches::search_model_weights::search_model_weights; +use elasticsearch_schema::searches::search_model_weights::{search_model_weights, SearchArgs}; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; @@ -115,12 +115,28 @@ pub async fn search_model_weights_handler( request: web::Json, server_state: web::Data>) -> Result { - let results = search_model_weights( - &server_state.elasticsearch, - &request.search_term, - request.ietf_language_subtag.as_deref(), - request.weight_type, - request.weight_category) + let maybe_weights_categories = request.weight_category + .map(|weight_category| { + let mut set = HashSet::new(); + set.insert(weight_category); + set + }); + + let maybe_weights_types = request.weight_type + .map(|weight_type| { + let mut set = HashSet::new(); + set.insert(weight_type); + set + }); + + let results = search_model_weights(SearchArgs { + search_term: &request.search_term, + maybe_creator_user_token: None, + maybe_ietf_primary_language_subtag: request.ietf_language_subtag.as_deref(), + maybe_weights_categories, + maybe_weights_types, + client: &server_state.elasticsearch, + }) .await .map_err(|err| { error!("Searching error: {:?}", err);