From ae2dee41e94a20033eec049267caee273d9995ab Mon Sep 17 00:00:00 2001
From: daxcay <daxtoncaylor@gmail.com>
Date: Tue, 2 Jul 2024 20:27:53 +0530
Subject: [PATCH] New Node

---
 __init__.py                     |  5 ++-
 classes/DATASET_TriggerWords.py | 61 +++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 classes/DATASET_TriggerWords.py

diff --git a/__init__.py b/__init__.py
index 5e91c94..83df9f9 100644
--- a/__init__.py
+++ b/__init__.py
@@ -30,6 +30,7 @@ from .classes.DATASET_LoadImage import N_CLASS_MAPPINGS as LoadImageMappings,  N
 from .classes.DATASET_SaveImage import N_CLASS_MAPPINGS as SaveImageMappings,  N_DISPLAY_NAME_MAPPINGS as SaveImageNameMappings
 from .classes.DATASET_TXTFileSaverBatch import N_CLASS_MAPPINGS as TXTFileSaverBatchMappings, N_DISPLAY_NAME_MAPPINGS as TXTFileSaverBatchNameMappings
 from .classes.DATASET_OpenAIChatImageBatch import N_CLASS_MAPPINGS as OpenAIChatImageBatchMappings,  N_DISPLAY_NAME_MAPPINGS as OpenAIChatImageBatchNameMappings
+from .classes.DATASET_TriggerWords import N_CLASS_MAPPINGS as TriggerWordsMappings,  N_DISPLAY_NAME_MAPPINGS as TriggerWordsNameMappings
 
 NODE_CLASS_MAPPINGS = {}
 NODE_CLASS_MAPPINGS.update(TXTFileSaverMappings)
@@ -44,6 +45,7 @@ NODE_CLASS_MAPPINGS.update(LoadImageMappings)
 NODE_CLASS_MAPPINGS.update(SaveImageMappings)
 NODE_CLASS_MAPPINGS.update(TXTFileSaverBatchMappings)
 NODE_CLASS_MAPPINGS.update(OpenAIChatImageBatchMappings)
+NODE_CLASS_MAPPINGS.update(TriggerWordsMappings)
 
 NODE_DISPLAY_NAME_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS.update(TXTFileSaverNameMappings)
@@ -56,7 +58,8 @@ NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatNameMappings)
 NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatImageNameMappings)
 NODE_DISPLAY_NAME_MAPPINGS.update(LoadImageNameMappings)
 NODE_DISPLAY_NAME_MAPPINGS.update(SaveImageNameMappings)
-NODE_DISPLAY_NAME_MAPPINGS.update(TXTFileSaverNameMappings)
+NODE_DISPLAY_NAME_MAPPINGS.update(TXTFileSaverBatchNameMappings)
 NODE_DISPLAY_NAME_MAPPINGS.update(OpenAIChatImageBatchNameMappings)
+NODE_DISPLAY_NAME_MAPPINGS.update(TriggerWordsNameMappings)
 
 WEB_DIRECTORY = "./web"
\ No newline at end of file
diff --git a/classes/DATASET_TriggerWords.py b/classes/DATASET_TriggerWords.py
new file mode 100644
index 0000000..3611f30
--- /dev/null
+++ b/classes/DATASET_TriggerWords.py
@@ -0,0 +1,61 @@
+def find_trigger_words(input_str, include_till_comma=False):
+    parts = [part.strip() for part in input_str.split(',')]    
+    trigger_words = []
+    for part in parts:
+        words = part.split()
+        for word in words:
+            if any(char.isdigit() for char in word):
+                if include_till_comma:
+                    trigger_words.append(part)
+                else:
+                    trigger_words.append(word)
+    
+    return ', '.join(trigger_words)
+
+class DATASET_TriggerWords:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "content": ("STRING",{"forceInput": True}),
+                "search": (['trigger_word_only','trigger_word_phrase'],),
+            },
+        }
+
+    INPUT_IS_LIST = True
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("Words",)
+    OUTPUT_IS_LIST = (True,)
+    FUNCTION = "FindIT"
+    OUTPUT_NODE = True
+
+    CATEGORY = "🔶DATASET🔶"
+
+    def FindIT(self, content, search):
+        try:
+
+            search = search[0]
+            words = []
+            
+            for sentence in content:
+                if search == "trigger_word_only": 
+                    words.append(find_trigger_words(sentence, False))
+                elif search == "trigger_word_phrase":
+                    words.append(find_trigger_words(sentence, True))            
+
+        except Exception as e:
+            print(f"Error saving: {e}")
+
+        return (words,)
+    
+N_CLASS_MAPPINGS = {
+    "DATASET_TriggerWords": DATASET_TriggerWords,
+}
+
+N_DISPLAY_NAME_MAPPINGS = {
+    "DATASET_TriggerWords": "DATASET_TriggerWords",
+}