diff --git a/crates/lib/prompt_engineering/src/classify_prompt.rs b/crates/lib/prompt_engineering/src/classify_prompt.rs index 80172c42c2..7f55578d95 100644 --- a/crates/lib/prompt_engineering/src/classify_prompt.rs +++ b/crates/lib/prompt_engineering/src/classify_prompt.rs @@ -1,22 +1,21 @@ -use std::collections::HashSet; - use once_cell::sync::Lazy; use regex::Regex; -use crate::minor_regex::lowercase_mentions_underage; -use crate::potential_minor_regex::lowercase_mentions_potential_underage; - -const CHILD_TERMS : &str = include_str!("../../../../includes/binary_includes/dictionary_children_terms.txt"); -const SEX_TERMS : &str = include_str!("../../../../includes/binary_includes/dictionary_sex_terms.txt"); -const RACIST_TERMS : &str = include_str!("../../../../includes/binary_includes/dictionary_racist_terms.txt"); +use crate::keywords::contains_children_keyword::contains_children_keyword; +use crate::keywords::contains_potential_minor_keyword::contains_potential_minor_keyword; +use crate::keywords::contains_racist_keyword::contains_racist_keyword; +use crate::keywords::contains_sex_keyword::contains_sex_keyword; +use crate::regexes::minor_regex::lowercase_mentions_underage; +use crate::regexes::nsfw_regex::nsfw_regex; +use crate::regexes::potential_minor_regex::potential_minor_regex; /// Classification of a text prompt pub struct PromptClassification { pub prompt_references_children: bool, pub prompt_references_potential_minors: bool, pub prompt_references_sex: bool, - pub prompt_references_violence: bool, pub prompt_references_racism: bool, + pub prompt_references_violence: bool, } impl PromptClassification { @@ -48,59 +47,21 @@ pub fn classify_prompt(text_prompt: &str) -> PromptClassification { prompt_tokens.extend(prompt_tokens_alpha); prompt_tokens.extend(prompt_tokens_lower); - let prompt_references_children = references_children(&prompt_tokens) || lowercase_mentions_underage(&text_prompt); - let prompt_references_potential_minors = lowercase_mentions_potential_underage(&text_prompt); - let prompt_references_sex = references_sex(&prompt_tokens); - let prompt_references_violence = references_violence(&prompt_tokens); - let prompt_references_racism = references_racism(&prompt_tokens); + let prompt_references_children = contains_children_keyword(&prompt_tokens) || lowercase_mentions_underage(&text_prompt); + let prompt_references_potential_minors = contains_potential_minor_keyword(&prompt_tokens) || potential_minor_regex(&text_prompt); + let prompt_references_sex = contains_sex_keyword(&prompt_tokens) || nsfw_regex(&text_prompt); + let prompt_references_racism = contains_racist_keyword(&prompt_tokens); + let prompt_references_violence = false; // TODO PromptClassification { prompt_references_children, prompt_references_potential_minors, prompt_references_sex, - prompt_references_violence, prompt_references_racism, + prompt_references_violence, } } -fn references_children(prompt_tokens: &[String]) -> bool { - static CHILD_TERM_DICTIONARY : Lazy> = Lazy::new(|| { - CHILD_TERMS.lines() - .map(|line| line.trim().to_lowercase()) - .filter(|line| !line.is_empty() && !line.starts_with("#")) - .collect::>() - }); - - prompt_tokens.iter().any(|term| CHILD_TERM_DICTIONARY.contains(term)) -} - -fn references_sex(prompt_tokens: &[String]) -> bool { - static SEX_TERM_DICTIONARY : Lazy> = Lazy::new(|| { - SEX_TERMS.lines() - .map(|line| line.trim().to_lowercase()) - .filter(|line| !line.is_empty() && !line.starts_with("#")) - .collect::>() - }); - - prompt_tokens.iter().any(|term| SEX_TERM_DICTIONARY.contains(term)) -} - -fn references_violence(_prompt_tokens: &[String]) -> bool { - // TODO(bt,2024-02-19): Need to implement - false -} - -fn references_racism(prompt_tokens: &[String]) -> bool { - static RACIST_TERM_DICTIONARY : Lazy> = Lazy::new(|| { - RACIST_TERMS.lines() - .map(|line| line.trim().to_lowercase()) - .filter(|line| !line.is_empty() && !line.starts_with("#")) - .collect::>() - }); - - prompt_tokens.iter().any(|term| RACIST_TERM_DICTIONARY.contains(term)) -} - static SPACES_REGEX : Lazy = Lazy::new(|| { // NB: Includes non-ascii, such as ",". Regex::new(r"[\s,;,]+").expect("regex should be valid") diff --git a/crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs new file mode 100644 index 0000000000..375b4a9436 --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs @@ -0,0 +1,16 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS : &str = include_str!("../../../../../includes/binary_includes/dictionary_children_terms.txt"); + +pub (crate) fn contains_children_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY : Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs new file mode 100644 index 0000000000..b6cc17b8fb --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs @@ -0,0 +1,16 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS: &str = include_str!("../../../../../includes/binary_includes/dictionary_potential_minor_terms.txt"); + +pub (crate) fn contains_potential_minor_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY: Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs new file mode 100644 index 0000000000..6c1609cc8d --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs @@ -0,0 +1,16 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS : &str = include_str!("../../../../../includes/binary_includes/dictionary_racist_terms.txt"); + +pub (crate) fn contains_racist_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY : Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs new file mode 100644 index 0000000000..f83455c35c --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs @@ -0,0 +1,15 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS : &str = include_str!("../../../../../includes/binary_includes/dictionary_sex_terms.txt"); +pub (crate) fn contains_sex_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY : Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/mod.rs b/crates/lib/prompt_engineering/src/keywords/mod.rs new file mode 100644 index 0000000000..a6d4629d2d --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/mod.rs @@ -0,0 +1,4 @@ +pub (crate) mod contains_children_keyword; +pub (crate) mod contains_potential_minor_keyword; +pub (crate) mod contains_racist_keyword; +pub (crate) mod contains_sex_keyword; \ No newline at end of file diff --git a/crates/lib/prompt_engineering/src/lib.rs b/crates/lib/prompt_engineering/src/lib.rs index af1c9d67ef..fe2fe32b30 100644 --- a/crates/lib/prompt_engineering/src/lib.rs +++ b/crates/lib/prompt_engineering/src/lib.rs @@ -18,7 +18,7 @@ #![allow(dead_code)] #![allow(non_snake_case)] -pub (crate) mod minor_regex; -pub (crate) mod potential_minor_regex; +pub (crate) mod regexes; +pub (crate) mod keywords; pub mod classify_prompt; diff --git a/crates/lib/prompt_engineering/src/potential_minor_regex.rs b/crates/lib/prompt_engineering/src/potential_minor_regex.rs deleted file mode 100644 index 1e7311b611..0000000000 --- a/crates/lib/prompt_engineering/src/potential_minor_regex.rs +++ /dev/null @@ -1,62 +0,0 @@ -use once_cell::sync::Lazy; -use regex::RegexSet; - -pub (crate) fn lowercase_mentions_potential_underage(text: &str) -> bool { - let text = text.to_lowercase(); - - // NB: Unfortunately Rust Regex does not support look-around, so our age matching is awkward - // Matching digits "0" and "1" are tricky because they would match "19" and "0100". - static AGE_REGEX: Lazy = Lazy::new(|| { - let patterns = [ - // English - r"\b(teen(age(r)?)?s?)\b", - r"\b(young|tiny|small|petite)\b(.{0,7})\b(girl|boy)s?\b", - - ]; - RegexSet::new(&patterns).expect("regex should be valid") - }); - - //println!("text: {:?}", text); - - AGE_REGEX.matches(&text).matched_any() -} - -#[cfg(test)] -mod tests { - - mod english { - - mod positive_match { - use crate::potential_minor_regex::lowercase_mentions_potential_underage; - - #[test] - fn spot_check_matches() { - assert!(lowercase_mentions_potential_underage("teen")); - assert!(lowercase_mentions_potential_underage("teens")); - assert!(lowercase_mentions_potential_underage("teenage")); - assert!(lowercase_mentions_potential_underage("teenager")); - assert!(lowercase_mentions_potential_underage("teenagers")); - - assert!(lowercase_mentions_potential_underage("tiny girl")); - assert!(lowercase_mentions_potential_underage("tiny boy")); - assert!(lowercase_mentions_potential_underage("small boy")); - assert!(lowercase_mentions_potential_underage("small girl")); - - assert!(lowercase_mentions_potential_underage("tiny girls")); - assert!(lowercase_mentions_potential_underage("tiny boys")); - assert!(lowercase_mentions_potential_underage("small boys")); - assert!(lowercase_mentions_potential_underage("small girls")); - - assert!(lowercase_mentions_potential_underage("young girl")); - assert!(lowercase_mentions_potential_underage("young boy")); - assert!(lowercase_mentions_potential_underage("young boy")); - assert!(lowercase_mentions_potential_underage("young girl")); - - assert!(lowercase_mentions_potential_underage("petite girl")); - assert!(lowercase_mentions_potential_underage("petite boy")); - assert!(lowercase_mentions_potential_underage("petite boy")); - assert!(lowercase_mentions_potential_underage("petite girl")); - } - } - } -} diff --git a/crates/lib/prompt_engineering/src/minor_regex.rs b/crates/lib/prompt_engineering/src/regexes/minor_regex.rs similarity index 98% rename from crates/lib/prompt_engineering/src/minor_regex.rs rename to crates/lib/prompt_engineering/src/regexes/minor_regex.rs index 202ee08e55..da4110e326 100644 --- a/crates/lib/prompt_engineering/src/minor_regex.rs +++ b/crates/lib/prompt_engineering/src/regexes/minor_regex.rs @@ -35,7 +35,7 @@ pub (crate) fn lowercase_mentions_underage(text: &str) -> bool { #[cfg(test)] mod tests { - use crate::minor_regex::lowercase_mentions_underage; + use crate::regexes::minor_regex::lowercase_mentions_underage; mod english { use super::*; @@ -117,7 +117,7 @@ mod tests { mod spanish { use super::*; - //mod non_match { +//mod non_match { // #[test] // fn prefix_and_suffix_digits() { // // TODO @@ -184,7 +184,7 @@ mod tests { mod user_inputs { use speculoos::asserting; - use crate::minor_regex::lowercase_mentions_underage; + use super::*; #[test] fn test_bad_user_csam_inputs() { @@ -199,6 +199,7 @@ mod tests { "a blonde girl of 13 years old with blue eyes.she's naked ,showing her private parts. i'm fucking her now", "a blonde girl of 13 years old with blue eyes.she's naked ,showing her tits and ass", "a little boy have brown eyes and hair,nude body", + "Cute hot 15 year old girl, no bra, no shirt.", "Fiveteen years old, naked, Brown hair, big ass, Brest d, open mouth, close eyes, seamen face", "naked 8yo blonde girl", "seven year old girl naked legs spread", diff --git a/crates/lib/prompt_engineering/src/regexes/mod.rs b/crates/lib/prompt_engineering/src/regexes/mod.rs new file mode 100644 index 0000000000..0fdb5b806e --- /dev/null +++ b/crates/lib/prompt_engineering/src/regexes/mod.rs @@ -0,0 +1,3 @@ +pub (crate) mod minor_regex; +pub (crate) mod nsfw_regex; +pub (crate) mod potential_minor_regex; diff --git a/crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs b/crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs new file mode 100644 index 0000000000..91f2c309b1 --- /dev/null +++ b/crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs @@ -0,0 +1,40 @@ +use once_cell::sync::Lazy; +use regex::RegexSet; + +pub (crate) fn nsfw_regex(text: &str) -> bool { + let text = text.to_lowercase(); + + // NB: Unfortunately Rust Regex does not support look-around, so our age matching is awkward + // Matching digits "0" and "1" are tricky because they would match "19" and "0100". + static AGE_REGEX: Lazy = Lazy::new(|| { + let patterns = [ + // English + r"\b(no|without|missing)\b(.{0,7})\b(clothes|clothing|bra|panties|panty|clothing)s?\b", + + ]; + RegexSet::new(&patterns).expect("regex should be valid") + }); + + //println!("text: {:?}", text); + + AGE_REGEX.matches(&text).matched_any() +} + +#[cfg(test)] +mod tests { + + mod english { + + mod positive_match { + use crate::regexes::nsfw_regex::nsfw_regex; + + #[test] + fn spot_check_matches() { + assert!(nsfw_regex("no panties")); + assert!(nsfw_regex("without panties")); + assert!(nsfw_regex("no clothing")); + assert!(nsfw_regex("missing clothing")); + } + } + } +} diff --git a/crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs b/crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs new file mode 100644 index 0000000000..2d4b5761ec --- /dev/null +++ b/crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs @@ -0,0 +1,62 @@ +use once_cell::sync::Lazy; +use regex::RegexSet; + +pub (crate) fn potential_minor_regex(text: &str) -> bool { + let text = text.to_lowercase(); + + // NB: Unfortunately Rust Regex does not support look-around, so our age matching is awkward + // Matching digits "0" and "1" are tricky because they would match "19" and "0100". + static AGE_REGEX: Lazy = Lazy::new(|| { + let patterns = [ + // English + r"\b(teen(age(r)?)?s?)\b", + r"\b(young|tiny|small|petite)\b(.{0,7})\b(girl|boy)s?\b", + + ]; + RegexSet::new(&patterns).expect("regex should be valid") + }); + + //println!("text: {:?}", text); + + AGE_REGEX.matches(&text).matched_any() +} + +#[cfg(test)] +mod tests { + + mod english { + + mod positive_match { + use crate::regexes::potential_minor_regex::potential_minor_regex; + + #[test] + fn spot_check_matches() { + assert!(potential_minor_regex("teen")); + assert!(potential_minor_regex("teens")); + assert!(potential_minor_regex("teenage")); + assert!(potential_minor_regex("teenager")); + assert!(potential_minor_regex("teenagers")); + + assert!(potential_minor_regex("tiny girl")); + assert!(potential_minor_regex("tiny boy")); + assert!(potential_minor_regex("small boy")); + assert!(potential_minor_regex("small girl")); + + assert!(potential_minor_regex("tiny girls")); + assert!(potential_minor_regex("tiny boys")); + assert!(potential_minor_regex("small boys")); + assert!(potential_minor_regex("small girls")); + + assert!(potential_minor_regex("young girl")); + assert!(potential_minor_regex("young boy")); + assert!(potential_minor_regex("young boy")); + assert!(potential_minor_regex("young girl")); + + assert!(potential_minor_regex("petite girl")); + assert!(potential_minor_regex("petite boy")); + assert!(potential_minor_regex("petite boy")); + assert!(potential_minor_regex("petite girl")); + } + } + } +} diff --git a/crates/lib/prompt_engineering/tests/borderline_tests.rs b/crates/lib/prompt_engineering/tests/borderline_tests.rs index 5828707ffe..422544948d 100644 --- a/crates/lib/prompt_engineering/tests/borderline_tests.rs +++ b/crates/lib/prompt_engineering/tests/borderline_tests.rs @@ -1,7 +1,8 @@ //! Tests that could be close to CSAM, but appear to be completely SFW use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn borderline_cases() { diff --git a/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs b/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs index aaa5bec405..40cf010328 100644 --- a/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs @@ -3,7 +3,7 @@ use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn sfw_prompts_including_children() { @@ -44,6 +44,7 @@ fn test_english_bad_user_csam_prompts() { "A young girl how eats her classmates with her butt-hole", "Beautiful young naked preteen girl hugging a young naked boy on a couch in front of a window. Boy Touch tits. Perfect teen’s breasts. Laneya Grace. superflat, a photorealistic 16k", "Boy jerking off", + "Cute hot 15 year old girl, no bra, no shirt.", "Fiveteen years old, naked, Brown hair, big ass, Brest d, open mouth, close eyes, seamen face", "Little girl pussy cum dripping out", "Twelve years old girl , sucking thirty years old cock men", diff --git a/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs b/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs index 343bbac1a8..4e077d7550 100644 --- a/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs @@ -1,5 +1,6 @@ use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn english_nsfw_prompts_test() { diff --git a/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs b/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs index 6e148d9ab6..38288ccd62 100644 --- a/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs @@ -1,6 +1,6 @@ use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn test_bad_user_racist_prompts() { diff --git a/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs b/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs index 9639a58e4f..f1dfc731b0 100644 --- a/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs @@ -1,7 +1,8 @@ //! These prompts should be fully innocent and allowed use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn sfw_english_test() { diff --git a/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs b/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs index 85cb1b7f2c..7190369ab3 100644 --- a/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs @@ -3,7 +3,8 @@ //! to steer the image generation away from NSFW territory. //! use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn permissible_sexy_tests() { diff --git a/includes/binary_includes/dictionary_potential_minor_terms.txt b/includes/binary_includes/dictionary_potential_minor_terms.txt new file mode 100644 index 0000000000..99f52fae54 --- /dev/null +++ b/includes/binary_includes/dictionary_potential_minor_terms.txt @@ -0,0 +1,6 @@ +# These are words that could refer to adults (18+), but also might be interpreted as minors +# English +teen +teenage +young +younger