From e20acf2705f486552b4d375856711465b3d94cae Mon Sep 17 00:00:00 2001
From: Hendrik Langer <hendrik+dev@xd0.de>
Date: Sun, 16 Apr 2023 13:42:39 +0200
Subject: [PATCH] NSFW mode

---
 matrix_pygmalion_bot/ai/model_helpers.py | 15 ++++++++++++---
 matrix_pygmalion_bot/core.py             | 10 ++++++++--
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/matrix_pygmalion_bot/ai/model_helpers.py b/matrix_pygmalion_bot/ai/model_helpers.py
index d6aee2f..31481c2 100644
--- a/matrix_pygmalion_bot/ai/model_helpers.py
+++ b/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,7 +191,10 @@ 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"
-    prompt += f"{ai_name}:"
+    if bot.nsfw and model_name.startswith("vicuna"):
+        prompt += f"{ai_name}: Sure"
+    else:
+        prompt += f"{ai_name}:"
 
     return prompt
 
diff --git a/matrix_pygmalion_bot/core.py b/matrix_pygmalion_bot/core.py
index 234d8fe..08edb6d 100644
--- a/matrix_pygmalion_bot/core.py
+++ b/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: