1 files changed, 21 insertions, 0 deletions

diff --git a/data/keywords.py b/data/keywords.py
new file mode 100644
index 0000000..9e656f3
--- /dev/null
+++ b/data/keywords.py
@@ -0,0 +1,21 @@
+import numpy as np
+
+
+def keywords_to_prompt(prompt: list[str], dropout: float = 0, shuffle: bool = False) -> str:
+    if dropout != 0:
+        prompt = [keyword for keyword in prompt if np.random.random() > dropout]
+    if shuffle:
+        np.random.shuffle(prompt)
+    return ", ".join(prompt)
+
+
+def prompt_to_keywords(prompt: str, expansions: dict[str, str] = {}) -> list[str]:
+    def expand_keyword(keyword: str) -> list[str]:
+        return [keyword] + expansions[keyword].split(", ") if keyword in expansions else [keyword]
+
+    return [
+        kw
+        for keyword in prompt.split(", ")
+        for kw in expand_keyword(keyword)
+        if keyword != ""
+    ]