Merge pull request #23 from yukiarimo/dev

Empty message error got fixed

index.py

@@ -1,4 +1,3 @@
-import base64
 from flask import Flask, request, jsonify
 from lib.generate import ChatGenerator, ChatHistoryManager
 from lib.vision import capture_image, create_image
@@ -85,9 +84,9 @@ def handle_message_request(self):
         chat_id = data.get('chat')
         speech = data.get('speech')
         text = data.get('text')
-        naked = data.get('naked')
+        template = data.get('template')
 
-        response = self.chat_generator.generate(chat_id, speech, text, naked)
+        response = self.chat_generator.generate(chat_id, speech, text, template)
 
         return jsonify({'response': response})
 

lib/generate.py

@@ -1,13 +1,9 @@
-import shutil
-import speech_recognition as sr
 import json
-import subprocess
-import subprocess
 import os
 import re
-import glob
 from transformers import pipeline
 from ctransformers import AutoModelForCausalLM
+from lib.history import ChatHistoryManager
 
 class ChatGenerator:
     def __init__(self, config):
@@ -29,10 +25,13 @@ def __init__(self, config):
         )
         self.classifier = pipeline("text-classification", model=f"{config['server']['agi_model_dir']}yuna-emotion")
 
-    def generate(self, chat_id, speech=False, text="", naked=False):
-        new_history = ''
+    def generate(self, chat_id, speech=False, text="", template=None):
+        if template == "dialog":
+            prompt_dir = os.path.join(self.config["server"]["prompts"] + 'dialog.txt')
+            with open(prompt_dir, 'r') as file:
+                prompt = file.read()
 
-        if naked==False:
+            history = ''
             with open(os.path.join(self.config["server"]["history"], chat_id), 'r') as file:
                 data = json.load(file)
                 history = ''
@@ -42,28 +41,45 @@ def generate(self, chat_id, speech=False, text="", naked=False):
                     if name and message:
                         history += f'{name}: {message}\n'
 
-            new_history = f"{history}Yuki: {text}\nYuna:"
-            new_history_crop = self.clearText(new_history)
-        elif naked:
-            prompts_dir = os.path.join(self.config["server"]["prompts"])
-            prompt_files = glob.glob(os.path.join(prompts_dir, '*.txt'))
+            history = f"{history}Yuki: {text}\nYuna:"
 
-            file_contents = {}  # Create a dictionary to store the file contents
+            # calculate the length of the prompt variable
+            prompt_length = len(self.model.tokenize(prompt))
 
-            for prompt_file in prompt_files:
-                with open(prompt_file, 'r') as file:
-                    file_content = file.read()
-                    file_content = file_content.replace('{user_msg}', text)
-                    variable_name = os.path.splitext(os.path.basename(prompt_file))[0]
-                    file_contents[variable_name] = file_content  # Store the file content in the dictionary
-
-            # inject prompt variable from naked into new_history
-            new_history = file_contents[naked]
+            # Calculate the maximum length for the history
+            max_length = self.config["ai"]["context_length"] - self.config["ai"]["max_new_tokens"]
 
-        new_history_crop = self.model.tokenize(new_history)
-        new_history_crop = new_history_crop[-((self.config["ai"]["context_length"] - self.config["ai"]["max_new_tokens"]) - 2):]
+            # Crop the history to fit within the max_length and prompt_length combined, counting from the end of the text
+            cropped_history = history[-(max_length - prompt_length):]
 
-        response = self.model(self.model.detokenize(new_history_crop), stream=False)
+            # replace string {user_msg} in the prompt with the history
+            response = prompt.replace('{user_msg}', cropped_history)
+
+            print(len(self.model.tokenize(response)), '<- response length')
+
+            # inject prompt variable from dialog into new_history
+        elif template != "dialog" and template != None:
+            prompt_dir = os.path.join(self.config["server"]["prompts"] + f'{template}.txt')
+            with open(prompt_dir, 'r') as file:
+                prompt = file.read()
+
+            # calculate the length of the prompt variable
+            prompt_length = len(self.model.tokenize(prompt))
+
+            # Calculate the maximum length for the history
+            max_length = self.config["ai"]["context_length"] - self.config["ai"]["max_new_tokens"]
+
+            # Crop the history to fit within the max_length and prompt_length combined, counting from the end of the text
+            cropped_history = text[:(max_length - prompt_length)]
+
+            # replace string {user_msg} in the prompt with the history
+            response = prompt.replace('{user_msg}', cropped_history)
+
+            print(len(self.model.tokenize(response)), '<- response length')
+        elif template == None:
+            print('template is none')
+
+        response = self.model(response, stream=False)
         response = self.clearText(str(response))
 
         if self.config["ai"]["emotions"]:
@@ -90,7 +106,7 @@ def generate(self, chat_id, speech=False, text="", naked=False):
         chat_history.append({"name": "Yuna", "message": response})
         ChatHistoryManager.save_chat_history(self, chat_history, chat_id)
 
-        if speech and "<tts>" in text:
+        if speech==True:
             self.generate_speech(response)
 
         return response
@@ -124,52 +140,6 @@ def remove_emojis(self, input_string):
 
         return cleaned_string
 
-class ChatHistoryManager:
-    def __init__(self, config):
-        self.config = config
-
-    def create_chat_history_file(self, chat_id):
-        history_starting_template = [{"name": "Yuki", "message": "Hi"}, {"name": "Yuna", "message": "Hello"}]
-        with open(os.path.join(self.config["server"]["history"], chat_id), 'w') as file:
-            json.dump(history_starting_template, file)
-
-    def delete_chat_history_file(self, chat_id):
-        os.remove(os.path.join(self.config["server"]["history"], chat_id))
-
-    def rename_chat_history_file(self, chat_id, new_chat_id):
-        os.rename(
-            os.path.join(self.config["server"]["history"], chat_id),
-            os.path.join(self.config["server"]["history"], new_chat_id)
-        )
-
-    def list_history_files(self):
-        history_files = [f for f in os.listdir(self.config["server"]["history"]) if os.path.isfile(os.path.join(self.config["server"]["history"], f))]
-
-        # Place main_history_file first and then sort alphabetically
-        history_files.sort(key=lambda x: (x != self.config["server"]["default_history_file"], x.lower()))
-
-        return history_files
-
-    def load_chat_history(self, chat):
-        print(self.config["server"]["history"], chat)
-        print(chat)
-        history_path = os.path.join(self.config["server"]["history"], chat)
-        if os.path.exists(history_path):
-            with open(history_path, 'r') as file:
-                return json.load(file)
-        else:
-            return []
-
-    def save_chat_history(self, chat_history, chat):
-        history_path = os.path.join(self.config["server"]["history"], chat)
-        with open(history_path, 'w') as file:
-            json.dump(chat_history, file)
-
-    def generate_speech(self, response):
-        subprocess.run(f'say "{response}" -o output', shell=True)
-        shutil.move("output.aiff", "static/audio/output.aiff")
-        subprocess.run(f"ffmpeg -y -i 'static/audio/output.aiff' -b:a 192K -f mp3 static/audio/output.mp3", shell=True)
-
 if __name__ == '__main__':
     with open("static/config.json", 'r') as file:
         config = json.load(file)

lib/history.py

@@ -0,0 +1,50 @@
+import json
+import os
+import shutil
+import subprocess
+
+class ChatHistoryManager:
+    def __init__(self, config):
+        self.config = config
+
+    def create_chat_history_file(self, chat_id):
+        history_starting_template = [{"name": "Yuki", "message": "Hi"}, {"name": "Yuna", "message": "Hello"}]
+        with open(os.path.join(self.config["server"]["history"], chat_id), 'w') as file:
+            json.dump(history_starting_template, file)
+
+    def delete_chat_history_file(self, chat_id):
+        os.remove(os.path.join(self.config["server"]["history"], chat_id))
+
+    def rename_chat_history_file(self, chat_id, new_chat_id):
+        os.rename(
+            os.path.join(self.config["server"]["history"], chat_id),
+            os.path.join(self.config["server"]["history"], new_chat_id)
+        )
+
+    def list_history_files(self):
+        history_files = [f for f in os.listdir(self.config["server"]["history"]) if os.path.isfile(os.path.join(self.config["server"]["history"], f))]
+
+        # Place main_history_file first and then sort alphabetically
+        history_files.sort(key=lambda x: (x != self.config["server"]["default_history_file"], x.lower()))
+
+        return history_files
+
+    def load_chat_history(self, chat):
+        print(self.config["server"]["history"], chat)
+        print(chat)
+        history_path = os.path.join(self.config["server"]["history"], chat)
+        if os.path.exists(history_path):
+            with open(history_path, 'r') as file:
+                return json.load(file)
+        else:
+            return []
+
+    def save_chat_history(self, chat_history, chat):
+        history_path = os.path.join(self.config["server"]["history"], chat)
+        with open(history_path, 'w') as file:
+            json.dump(chat_history, file)
+
+    def generate_speech(self, response):
+        subprocess.run(f'say "{response}" -o output', shell=True)
+        shutil.move("output.aiff", "static/audio/output.aiff")
+        subprocess.run(f"ffmpeg -y -i 'static/audio/output.aiff' -b:a 192K -f mp3 static/audio/output.mp3", shell=True)

static/config.json

@@ -4,7 +4,7 @@
         "emotions": false,
         "art": false,
         "max_new_tokens": 128,
-        "context_length": 512,
+        "context_length": 1024,
         "temperature": 0.7,
         "repetition_penalty": 1.2,
         "last_n_tokens": 128,