Update app.py

app.py

@@ -37,6 +37,7 @@ def generate(
     # Clear the stop event before starting a new generation
     stop_event.clear()
 
+    # Append the user's message to the conversation history
     conversation = chat_history.copy()
     conversation.append({"role": "user", "content": message})
 
@@ -46,6 +47,7 @@ def generate(
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
 
+    # Create a streamer to get the generated response
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         {"input_ids": input_ids},
@@ -53,6 +55,7 @@ def generate(
         max_new_tokens=max_new_tokens,
     )
     
+    # Run generation in a background thread
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
@@ -63,6 +66,11 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
+    # After generation, append the assistant's response to the chat history
+    assistant_response = "".join(outputs)
+    chat_history.append({"role": "assistant", "content": assistant_response})
+
+
 # Define a function to stop the generation
 def stop_generation():
     stop_event.set()