Update app.py

app.py

@@ -19,8 +19,15 @@ def generate_response(input_text):
     # Encode the new user input, add end of string token
     new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt').to(device)
 
-    # Append the new user input tokens to the chat history
-    bot_input_ids = torch.cat([torch.tensor(st.session_state['history']).to(device), new_user_input_ids], dim=-1) if st.session_state['history'] else new_user_input_ids
+    # If there is conversation history, append the new input to it
+    if st.session_state['history']:
+        # Convert history to a 2D tensor (batch_size x seq_len)
+        history_tensor = torch.tensor(st.session_state['history']).unsqueeze(0).to(device)
+        # Concatenate history with the new input
+        bot_input_ids = torch.cat([history_tensor, new_user_input_ids], dim=-1)
+    else:
+        # If no history, just use the new user input
+        bot_input_ids = new_user_input_ids
 
     # Generate a response from the model
     chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id, top_k=50, top_p=0.95, temperature=0.7)
@@ -29,7 +36,7 @@ def generate_response(input_text):
     chat_history_ids = chat_history_ids[:, bot_input_ids.shape[-1]:]  # only take the latest generated tokens
     bot_output = tokenizer.decode(chat_history_ids[0], skip_special_tokens=True)
 
-    # Update session state history with the new tokens
+    # Update session state history with the new tokens (flattened)
     st.session_state['history'] = chat_history_ids[0].tolist()
 
     return bot_output