NSFW mode

matrix_pygmalion_bot/ai/model_helpers.py

@@ -54,11 +54,17 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history, model_name: str
         prompt = ""
     elif model_name.startswith("vicuna"):
         prompt = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
-        prompt += f"### Instruction:\nGiven the following character description and scenario, write a script for a dialogue between the human user {bot.user_name} and the fictional AI assistant {bot.name}. Play the role of the character {bot.name}\n\n" # ToDo
+        prompt += f"### Instruction:\nGiven the following character description and scenario, write a script for a dialogue between the human user {bot.user_name} and the fictional AI assistant {bot.name}. Play the role of the character {bot.name}.\n" # ToDo
+        if bot.nsfw:
+            prompt += f"NSFW content allowed.\n" # ToDo
+        prompt += f"\n"
         prompt += "### Input:\n"
     elif model_name.startswith("alpaca"):
         prompt = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
-        prompt += f"### Instruction:\nGiven the following character description and scenario, write a script for a dialogue between the human user {bot.user_name} and the fictional AI assistant {bot.name}. Play the role of the character {bot.name}\n\n"
+        prompt += f"### Instruction:\nGiven the following character description and scenario, write a script for a dialogue between the human user {bot.user_name} and the fictional AI assistant {bot.name}. Play the role of the character {bot.name}.\n"
+        if bot.nsfw:
+            prompt += f"NSFW content allowed.\n" # ToDo
+        prompt += f"\n"
         prompt += "### Input:\n"
     else:
         prompt = ""
@@ -185,6 +191,9 @@ async def get_full_prompt(simple_prompt: str, bot, chat_history, model_name: str
         chat_history.setFastForward(True)
 
     prompt += f"{user_name}: {simple_prompt}\n"
+    if bot.nsfw and model_name.startswith("vicuna"):
+        prompt += f"{ai_name}: Sure"
+    else:
         prompt += f"{ai_name}:"
 
     return prompt

matrix_pygmalion_bot/core.py

@@ -325,6 +325,7 @@ class ChatBot(object):
         self.scenario = None
         self.greeting = None
         self.example_dialogue = []
+        self.nsfw = False
         self.keywords = []
         self.extra_info = {"persona": [], "scenario": [], "example_dialogue": []}
         self.temperature = 0.90
@@ -338,13 +339,14 @@ class ChatBot(object):
         if STORE_PATH and not os.path.isdir(STORE_PATH):
             os.mkdir(STORE_PATH)
 
-    def character_init(self, name, persona, scenario, greeting, example_dialogue=[]):
+    def character_init(self, name, persona, scenario, greeting, example_dialogue=[], nsfw=False):
         self.name = name
         self.persona = persona
         self.scenario = scenario
         self.greeting = greeting
         self.example_dialogue = example_dialogue
         self.chat_history = BotChatHistory(self.name)
+        self.nsfw = nsfw
 
     def get_persona(self):
         return ' '.join([self.persona, ' '.join(self.extra_info['persona'])])
@@ -513,7 +515,11 @@ async def main() -> None:
             example_dialogue = json.loads(config[section]['example_dialogue'])
         else:
             example_dialogue = []
-        bot.character_init(botname, config[section]['persona'].replace("\\n", "\n"), config[section]['scenario'].replace("\\n", "\n"), config[section]['greeting'].replace("\\n", "\n"), example_dialogue)
+        if config.has_option(section, 'nsfw'):
+            nsfw = config[section]['nsfw']
+        else:
+            nsfw = False
+        bot.character_init(botname, config[section]['persona'].replace("\\n", "\n"), config[section]['scenario'].replace("\\n", "\n"), config[section]['greeting'].replace("\\n", "\n"), example_dialogue, nsfw)
         if config.has_option(section, 'keywords'):
             bot.keywords = json.loads(config[section]['keywords'])
         else: