1 files changed, 5 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index dbfe58c..5c0299e 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -648,7 +648,11 @@ def main():
         weight_dtype = torch.bfloat16
 
     def keyword_filter(item: CSVDataItem):
-        return any(keyword in item.prompt for keyword in args.placeholder_token)
+        return any(
+            keyword in part
+            for keyword in args.placeholder_token
+            for part in item.prompt
+        )
 
     def collate_fn(examples):
         prompts = [example["prompts"] for example in examples]