✨ feat(playground): Add random_prompt method to RandomPromptGenerator

LlmKira · Feb 7, 2024 · ef4c96a · ef4c96a
1 parent b3e31d2
commit ef4c96a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -64,7 +64,7 @@ loop.run_until_complete(main())
 ```python
 from novelai_python.utils.random_prompt import RandomPromptGenerator
 
-s = RandomPromptGenerator(nsfw_enabled=False).generate()
+s = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
 print(s)
 ```
 

diff --git a/playground/random_prompt.py b/playground/random_prompt.py
@@ -14,5 +14,5 @@
 print(gen.get_weighted_choice([['m', 30], ['f', 50], ['o', 10]], ['m']))
 print("====")
 for i in range(200):
-    s = RandomPromptGenerator(nsfw_enabled=True).generate()
+    s = RandomPromptGenerator(nsfw_enabled=True).random_prompt()
     print(s)
diff --git a/src/novelai_python/utils/random_prompt/__init__.py b/src/novelai_python/utils/random_prompt/__init__.py
@@ -199,16 +199,21 @@ def random_prompt(self, *,
                       enable_moods: bool = True,
                       enable_character: bool = True,
                       enable_identity: bool = False,
+                      must_appear=None,
                       ):
+        if must_appear is None:
+            must_appear = []
         tags = []
+        # 必须出现的标签
+        tags.extend(must_appear)
         if self.nsfw_enabled:
             tags.append('nsfw')
             if random.random() < 0.1:
                 tags.append('explicit')
                 tags.append('lewd')
         irs = self.get_weighted_choice([[1, 70], [2, 20], [3, 7], [0, 5]], tags)
         if self.nsfw_enabled:
-            irs = self.get_weighted_choice([[1, 35], [2, 20], [3, 7]], tags)
+            irs = self.get_weighted_choice([[1, 40], [2, 20], [3, 7]], tags)
         if irs == 0:
             tags.append('no humans')
             if random.random() < 0.3: