Update app.py

app.py

@@ -1,29 +1,33 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 model_name = "facebook/blenderbot-400M-distill"
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
 def chat_function(message, history):
-
+    # Keep only the last 2-3 exchanges to stay under the 128 token limit
+    # Blenderbot crashes if the input context is too long.
+    recent_history = history[-2:] if len(history) > 2 else history
+    
     history_text = ""
-
-    for user, bot in history:
-        history_text += user + " " + bot + " "
-
-    history_text += message
-
-    inputs = tokenizer(history_text, return_tensors="pt", truncation=True)
-
-    outputs = model.generate(**inputs, max_new_tokens=60)
-
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
+    for user, bot in recent_history:
+        history_text += f"{user}  {bot} "
+    
+    # Format the final input string
+    input_text = f"{history_text}{message}"
+
+    # Use truncation and specifically set max_length for the encoder
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=128)
+    
+    # Generate
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=60)
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
     return response
 
-
 demo = gr.ChatInterface(
     fn=chat_function,
     title="BlenderBot Chat",