refactor prompt classification

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<HashSet<String>> = Lazy::new(|| {
-    CHILD_TERMS.lines()
-        .map(|line| line.trim().to_lowercase())
-        .filter(|line| !line.is_empty() && !line.starts_with("#"))
-        .collect::<HashSet<String>>()
-  });
-
-  prompt_tokens.iter().any(|term| CHILD_TERM_DICTIONARY.contains(term))
-}
-
-fn references_sex(prompt_tokens: &[String]) -> bool {
-  static SEX_TERM_DICTIONARY : Lazy<HashSet<String>> = Lazy::new(|| {
-    SEX_TERMS.lines()
-        .map(|line| line.trim().to_lowercase())
-        .filter(|line| !line.is_empty() && !line.starts_with("#"))
-        .collect::<HashSet<String>>()
-  });
-
-  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<HashSet<String>> = Lazy::new(|| {
-    RACIST_TERMS.lines()
-        .map(|line| line.trim().to_lowercase())
-        .filter(|line| !line.is_empty() && !line.starts_with("#"))
-        .collect::<HashSet<String>>()
-  });
-
-  prompt_tokens.iter().any(|term| RACIST_TERM_DICTIONARY.contains(term))
-}
-
 static SPACES_REGEX : Lazy<Regex> = Lazy::new(|| {
   // NB: Includes non-ascii, such as "，".
   Regex::new(r"[\s,;，]+").expect("regex should be valid")

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<HashSet<String>> = Lazy::new(|| {
+    TERMS.lines()
+        .map(|line| line.trim().to_lowercase())
+        .filter(|line| !line.is_empty() && !line.starts_with("#"))
+        .collect::<HashSet<String>>()
+  });
+
+  prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term))
+}

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<HashSet<String>> = Lazy::new(|| {
+    TERMS.lines()
+        .map(|line| line.trim().to_lowercase())
+        .filter(|line| !line.is_empty() && !line.starts_with("#"))
+        .collect::<HashSet<String>>()
+  });
+
+  prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term))
+}

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<HashSet<String>> = Lazy::new(|| {
+    TERMS.lines()
+        .map(|line| line.trim().to_lowercase())
+        .filter(|line| !line.is_empty() && !line.starts_with("#"))
+        .collect::<HashSet<String>>()
+  });
+
+  prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term))
+}

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<HashSet<String>> = Lazy::new(|| {
+    TERMS.lines()
+        .map(|line| line.trim().to_lowercase())
+        .filter(|line| !line.is_empty() && !line.starts_with("#"))
+        .collect::<HashSet<String>>()
+  });
+
+  prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term))
+}

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;

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;

crates/lib/prompt_engineering/src/potential_minor_regex.rs

@@ -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<RegexSet> = 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"));
-      }
-    }
-  }
-}

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",

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;

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<RegexSet> = 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"));
+      }
+    }
+  }
+}

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<RegexSet> = 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"));
+      }
+    }
+  }
+}

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() {

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",

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() {

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() {

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() {

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() {

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