Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
 from langchain.memory import ConversationBufferMemory
 
@@ -7,6 +7,19 @@ from langchain.memory import ConversationBufferMemory
 tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
 model = GPT2LMHeadModel.from_pretrained("distilgpt2")
 
+# Load summarization model (e.g., T5-small)
+summarizer_tokenizer = AutoTokenizer.from_pretrained("t5-small")
+summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("t5-small").to(device)
+
+def summarize_history(history):
+    input_ids = summarizer_tokenizer.encode(
+        "summarize: " + history,
+        return_tensors="pt"
+    ).to(device)
+    summary_ids = summarizer_model.generate(input_ids, max_length=50, min_length=25, length_penalty=5., num_beams=2)
+    summary = summarizer_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
+
 # Move model to device (GPU if available)
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 model.to(device)
@@ -16,17 +29,21 @@ memory = ConversationBufferMemory()
 
 # Define the chatbot function with memory
 def chat_with_distilgpt2(input_text):
-    # Retrieve conversation history and append the current user input
+    # Retrieve conversation history
     conversation_history = memory.load_memory_variables({})['history']
     
-    # Combine the history with the current user input
+    # Summarize if history exceeds certain length
+    if len(conversation_history.split()) > 200:
+        conversation_history = summarize_history(conversation_history)
+    
+    # Combine the (possibly summarized) history with the current user input
     full_input = f"{conversation_history}\nUser: {input_text}\nAssistant:"
     
     # Tokenize the input and convert to tensor
     input_ids = tokenizer.encode(full_input, return_tensors="pt").to(device)
     
     # Generate the response using the model
-    outputs = model.generate(input_ids, max_length=400, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+    outputs = model.generate(input_ids, max_length=150, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
     
     # Decode the model output
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)