Update app.py

app.py

@@ -19,46 +19,42 @@ def save_database(database):
         json.dump(database, file)
 
 def format_prompt(message, system, history):
-    # Format prompt according to the new template
     prompt = f"SYSTEM: {system}\n<|endofsystem|>\n"
     
     for entry in history:
-        if len(entry) == 2:  # Check if the history contains exactly two items
+        if len(entry) == 2:
             user_prompt, bot_response = entry
             prompt += f"USER: {user_prompt}\n\n\nASSISTANT: {bot_response}<|endoftext|>\n"
-        else:
-            # If the history doesn't match the expected format, handle it gracefully
-            continue
     
     prompt += f"USER: {message}\n\n\nASSISTANT:"
     return prompt
 
-def generate(
-    prompt, system, history, temperature=0.9, max_new_tokens=4096, top_p=0.9, repetition_penalty=1.2,
-):
-    database = load_database()  # Load the database
+def generate(prompt, system, history, temperature=0.9, max_new_tokens=4096, top_p=0.9, repetition_penalty=1.2):
+    database = load_database()
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
     top_p = float(top_p)
 
     formatted_prompt = format_prompt(prompt, system, history)
-    response_text = "We are sorry but Quble doesnt know how to answer."  # Default response
+    response_text = "We are sorry but Quble doesn't know how to answer."
 
     if formatted_prompt in database:
         response_text = database[formatted_prompt]
     else:
         # Use the pipeline to generate the response
         try:
-            response = pipe(formatted_prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)[0]["generated_text"]
-            response_text = response.split("ASSISTANT:")[1].strip()  # Extract the assistant's response
-            database[formatted_prompt] = response_text
-            save_database(database)  # Save the updated database
+            # Stream the response
+            for response in pipe(formatted_prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty, return_full_text=False, streaming=True):
+                assistant_response = response["generated_text"].split("ASSISTANT:")[-1].strip()
+                yield assistant_response
+
+            # Save the generated response to the database after the stream
+            database[formatted_prompt] = assistant_response
+            save_database(database)
         except Exception as e:
             print(f"Error generating response: {e}")
 
-    yield response_text
-
 customCSS = """
 #component-7 { # this is the default element ID of the chat component
   height: 1600px; # adjust the height as needed
@@ -66,7 +62,7 @@ customCSS = """
 }
 """
 
-additional_inputs=[
+additional_inputs = [
     gr.Textbox(
         label="System prompt",
         value="You are a helpful assistant, with no access to external functions.",