lock input when generating response

rbourgeat · Jul 16, 2023 · 6296d5b · 6296d5b
1 parent 155833c
commit 6296d5b
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 43 deletions.
diff --git a/app/.DS_Store b/app/.DS_Store
diff --git a/app/run.py b/app/run.py
@@ -31,7 +31,6 @@
 ##############################################################
 
 PROCESS = None
-CUSTOM = False
 MODEL_7B = "llama.cpp/models/WizardLM-7B-V1.0-Uncensored/ggml-model-q4_0.bin"
 MODEL_13B = "llama.cpp/models/WizardLM-13B-V1.0-Uncensored/ggml-model-q4_0.bin"
 MODEL_33B = "llama.cpp/models/WizardLM-33B-V1.0-Uncensored/ggml-model-q4_0.bin"
@@ -41,6 +40,7 @@
 output_queue = Queue()
 input_queue = Queue()
 is_generating_image = False  # pylint: disable=invalid-name
+custom = False  # pylint: disable=invalid-name
 
 nlp = spacy.load("en_core_web_sm")
 
@@ -67,9 +67,10 @@
     pipe = pipe.to("cuda")
 pipe.enable_attention_slicing()  # Recommended if your computer has < 64 GB of RAM
 
-werkzeug_logger = logging.getLogger('werkzeug')
+werkzeug_logger = logging.getLogger("werkzeug")
 werkzeug_logger.setLevel(logging.ERROR)
 
+
 @app.route("/")
 def index():
     """
@@ -91,10 +92,10 @@ def execute():
     """
     global system  # pylint: disable=invalid-name, global-variable-not-assigned
     mode = request.form["mode"]
-    if (mode == "custom"):
-        CUSTOM = True
+    if mode == "custom":  # pylint: disable=simplifiable-if-statement
+        custom = True  # pylint: disable=redefined-outer-name
     else:
-        CUSTOM = False
+        custom = False
 
     model = "WizardLM-7B-V1.0-Uncensored"
     if os.path.exists(MODEL_13B) and not LIGHT_MODE:
@@ -105,15 +106,15 @@ def execute():
 
     command = ""
     if system == "Darwin":
-        if CUSTOM:
+        if custom:
             command = f'./llama.cpp/main -m llama.cpp/models/{model}/ggml-model-q4_0.bin \
                 -ngl 1 --repeat_penalty 1.1 --color --interactive-first \
                 -f app/prompts/CustomRolePlay.txt -r "USER: "'
         else:
             command = f'./llama.cpp/main -m llama.cpp/models/{model}/ggml-model-q4_0.bin \
                 -ngl 1 --repeat_penalty 1.1 --color -i -f app/prompts/RolePlay.txt -r "USER: "'
     elif system == "Linux":
-        if CUSTOM:
+        if custom:
             command = f'./llama.cpp/main -m llama.cpp/models/{model}/ggml-model-q4_0.bin \
                 --repeat_penalty 1.1 --color --interactive-first \
                 -f app/prompts/RolePlay.txt -r "USER: "'
@@ -169,7 +170,7 @@ def get_output():
     try:
         output = output_queue.get(timeout=1.0)
         if "USER:" in output:
-            output = ""
+            output = "[EOS]"
     except Empty:
         output = ""
     return jsonify(output=output)

diff --git a/app/static/js/script.js b/app/static/js/script.js
@@ -3,21 +3,36 @@ var isRunning = false;
 var rpStarted = false;
 var modelLoading = 0;
 var prompt = "";
+var mode = "";
 
-function startExecution(mode) {
+function startExecution() {
     if (isRunning) return;
     isRunning = true;
     $.post('/execute', {mode: mode} , function(data) {
         pollOutput();
     });
-    document.getElementById("start-button-random").remove();
-    document.getElementById("start-button-custom").remove();
+    document.getElementById("button-container").remove();
+    $('#input-command').prop('disabled', true);
+    $('#input-command').val('Generating response...');
+    $('#send-button').prop('disabled', true);
+    $('#send-button').val('🔒');
 }
 
 function pollOutput() {
     $.get('/get_output', function(data) {
         var output = data.output;
-        if (output && rpStarted) {
+        if (!rpStarted) {
+            var progressBar = document.querySelector('.progress');
+            modelLoading += 0.07;
+            progressBar.style.width = String(modelLoading) + '%';
+        }
+        if (isRunning) {
+            setTimeout(pollOutput, pollingInterval);
+        }
+        if (output.includes("===")) {
+            rpStarted = true;
+        }
+        if (output && rpStarted && output !== "[EOS]") {
             prompt = prompt + output;
             if (prompt.length > 300)
                 prompt = prompt.slice(output.length);
@@ -29,35 +44,31 @@ function pollOutput() {
                 progressBar.style.width = String(modelLoading) + '%';
             }
         }
-        if (!rpStarted) {
-            var progressBar = document.querySelector('.progress');
-            modelLoading += 0.07;
-            progressBar.style.width = String(modelLoading) + '%';
-        }
-        if (isRunning) {
-            setTimeout(pollOutput, pollingInterval);
-        }
-        if (output.includes("===")) {
-            rpStarted = true;
+        if (output === "[EOS]" && rpStarted) {
+            $('#input-command').prop('disabled', false);
+            $('#input-command').val('');
+            $('#send-button').prop('disabled', false);
+            $('#send-button').val('Send');
+            $.post('/generate_image', { prompt: prompt }, function(data) {
+                if (data.error)
+                    return;
+                var imageContainer = document.getElementById('image-container');
+                var image = document.createElement('img');
+                image.src = "/images/" + data.file_name;
+                imageContainer.innerHTML = '';
+                imageContainer.appendChild(image);
+            })
         }
     });
 }
 
-
 function sendInput() {
     var input = $('#input-command').val().trim();
-    $('#input-command').val('');
+    $('#input-command').val('Generating response...');
+    $('#input-command').prop('disabled', true);
+    $('#send-button').prop('disabled', true);
+    $('#send-button').val('🔒');
     $.post('/send_input', {input: input});
-
-    $.post('/generate_image', { prompt: prompt != "" ? prompt : input }, function(data) {
-        if (data.error)
-            return;
-        var imageContainer = document.getElementById('image-container');
-        var image = document.createElement('img');
-        image.src = "/images/" + data.file_name;
-        imageContainer.innerHTML = '';
-        imageContainer.appendChild(image);
-    })
 }
 
 function reloadPage() {
@@ -76,13 +87,15 @@ $(document).ready(function() {
     $('#start-button-random').click(function() {
         $(this).hide();
         $('#input-form').show();
-        startExecution("random");
+        mode = "random";
+        startExecution();
     });
 
     $('#start-button-custom').click(function() {
         $(this).hide();
         $('#input-form').show();
-        startExecution("custom");
+        mode = "custom";
+        startExecution();
     });
 
     $('#input-form').submit(function(event) {

diff --git a/app/static/styles.css b/app/static/styles.css
@@ -57,13 +57,14 @@ body {
     overflow-wrap: break-word;
 }
 
-#input-form, #download-form {
-    width: 100%;
-    text-align: center;
+#input-form {
+    display: flex;
+    justify-content: center;
+    align-items: center;
 }
 
 #input-command {
-    width: 70%;
+    width: 400px;
     padding: 10px;
     border: none;
     border-radius: 25px;
@@ -99,7 +100,7 @@ button {
     color: white;
     border-radius: 25px;
     cursor: pointer;
-    margin-bottom: 20px;
+    margin: 0 20px;
 }
 
 #reload-button {
@@ -172,3 +173,9 @@ img {
     padding: 40px;
     width: 120%;
 }
+
+#button-container {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+}
diff --git a/app/templates/index.html b/app/templates/index.html
@@ -16,8 +16,10 @@ <h1>📖 LLM RP</h1>
             </div>
         </div>
         <div id="status-container"></div>
-        <button id="start-button-random">🎲 Start a random story</button>
-        <button id="start-button-custom">📝 Start a custom story</button>
+        <div id="button-container">
+            <button id="start-button-random">🎲 Start a random story</button>
+            <button id="start-button-custom">📝 Start a custom story</button>
+        </div>
         <form id="input-form">
             <input type="text" id="input-command" placeholder="Write action here" required autofocus autocomplete="off">
             <input type="submit" id="send-button" value="Send">