Adding clear history button and using last 100 tokens of history

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Config
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import torch
 from langchain.memory import ConversationBufferMemory
 
@@ -16,21 +16,39 @@ model.to(device)
 # Set up conversational memory using LangChain's ConversationBufferMemory
 memory = ConversationBufferMemory()
 
+# Function to truncate tokens to the last 100 tokens
+def truncate_history_to_100_tokens(history, tokenizer, max_tokens=100):
+    # Tokenize the history
+    tokenized_history = tokenizer.encode(history)
+    
+    # Truncate to the last 100 tokens if necessary
+    if len(tokenized_history) > max_tokens:
+        tokenized_history = tokenized_history[-max_tokens:]
+    
+    return tokenized_history
+
 # Define the chatbot function with memory and additional parameters
 def chat_with_dialogpt(input_text, temperature, top_p, top_k):
     # Retrieve conversation history
     conversation_history = memory.load_memory_variables({})['history']
     
     # Combine the (possibly summarized) history with the current user input
-    no_memory_input = f"Question: {input_text}\nAnswer:"
-    
-    # Tokenize the input and convert to tensor
-    input_ids = tokenizer.encode(no_memory_input, return_tensors="pt").to(device)
-    
+    full_history = conversation_history + f"\nYou: {input_text}\nBot:"
+
+    # Truncate history to the most recent 100 tokens
+    truncated_input_ids = truncate_history_to_100_tokens(full_history, tokenizer)
+
+    # Tokenize the user input and append to truncated history
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    truncated_input_ids_tensor = torch.tensor([truncated_input_ids]).to(device)
+
+    # Concatenate truncated history with the new input
+    final_input_ids = torch.cat((truncated_input_ids_tensor, input_ids), dim=1)
+
     # Generate the response using the model with adjusted parameters
     outputs = model.generate(
-        input_ids,
-        max_length=input_ids.shape[1] + 50,  # Limit total length
+        final_input_ids,
+        max_length=final_input_ids.shape[1] + 50,  # Limit total length
         max_new_tokens=15,
         num_return_sequences=1,
         no_repeat_ngram_size=3,
@@ -50,10 +68,15 @@ def chat_with_dialogpt(input_text, temperature, top_p, top_k):
     memory.save_context({"input": input_text}, {"output": response})
     
     # Format the chat history for display
-    chat_history = conversation_history + f"\nYou: {input_text}\nBot: {response}\n"
+    chat_history = full_history + f"\nBot: {response}\n"
     
     return chat_history
 
+# Function to clear the chat history
+def clear_history():
+    memory.clear()  # Clear the memory object
+    return ""  # Return empty string to reset the chat display
+
 # Set up the Gradio interface with the input box below the output box
 with gr.Blocks() as interface:
     chatbot_output = gr.Textbox(label="Conversation", lines=15, placeholder="Chat history will appear here...", interactive=False)
@@ -79,10 +102,12 @@ with gr.Blocks() as interface:
                       inputs=[user_input, chatbot_output, temperature_slider, top_p_slider, top_k_slider], 
                       outputs=[chatbot_output, user_input])
 
+    # Add a clear history button
+    clear_button = gr.Button("Clear History")
+    clear_button.click(fn=clear_history, outputs=[chatbot_output])
+
     # Layout for sliders and chatbot UI
     gr.Row([temperature_slider, top_p_slider, top_k_slider])
     
 # Launch the Gradio app
 interface.launch()
-
-