Update app.py

app.py

@@ -5,55 +5,71 @@ import torch
 # Load model and tokenizer
 model_name = "microsoft/DialoGPT-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# Move model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
 def respond(message, chat_history, chat_history_ids):
-    # Encode user input
-    new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
-    
-    # Append to chat history
-    if chat_history_ids is not None:
-        input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
-    else:
-        input_ids = new_input_ids
-    
-    # Generate response
-    chat_history_ids = model.generate(
-        input_ids,
-        max_length=1000,
-        pad_token_id=tokenizer.eos_token_id,
-        no_repeat_ngram_size=3,
-        do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        temperature=0.8
-    )
-    
-    # Decode response
-    response = tokenizer.decode(
-        chat_history_ids[:, input_ids.shape[-1]:][0], 
-        skip_special_tokens=True
-    )
+    if not message.strip():
+        return "", chat_history or [], chat_history_ids, "Please enter a message."
+
+    if chat_history is None:
+        chat_history = []
+
+    new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt").to(device)
     
-    # Update conversation history
-    chat_history.append((message, response))
+    input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1) if chat_history_ids is not None else new_input_ids
     
-    return "", chat_history, chat_history_ids
+    try:
+        chat_history_ids = model.generate(
+            input_ids,
+            max_length=200,
+            pad_token_id=tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.8
+        )
+        
+        response = tokenizer.decode(
+            chat_history_ids[:, input_ids.shape[-1]:][0], 
+            skip_special_tokens=True
+        )
+        
+        chat_history.append((message, response))
+        
+        if len(chat_history) > 10:
+            chat_history = chat_history[-10:]
+            history_text = "".join([msg + resp + tokenizer.eos_token for msg, resp in chat_history])
+            chat_history_ids = tokenizer.encode(history_text, return_tensors="pt").to(device)
+        
+        return "", chat_history, chat_history_ids, None
+    except Exception as e:
+        return "", chat_history, chat_history_ids, f"Error: {str(e)}"
+
+def clear_history():
+    return [], None, None
 
 with gr.Blocks() as demo:
-    # Store model's conversation history
     state = gr.State()
-    
     gr.Markdown("## DialoGPT Chatbot")
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your Message")
+    msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")
     clear = gr.Button("Clear History")
+    error = gr.Textbox(label="Error", interactive=False, visible=False)
 
     msg.submit(
         respond,
-        [msg, chatbot, state],
-        [msg, chatbot, state]
+        inputs=[msg, chatbot, state],
+        outputs=[msg, chatbot, state, error]
     )
-    clear.click(lambda: (None, None), outputs=[chatbot, state], queue=False)
-
-demo.launch()
+    clear.click(
+        fn=clear_history,
+        inputs=None,
+        outputs=[chatbot, state, error],
+        queue=False
+    )