Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 # Create a class for the session state
 class SessionState:
@@ -33,10 +34,15 @@ if input_text:
     history_string = "\n".join(message for role, message in session_state.conversation_history)
 
     # Tokenize the input text and history
-    inputs = tokenizer.encode_plus(history_string, input_text, return_tensors="pt")
+    inputs = tokenizer.encode_plus(history_string, return_tensors="pt")
+    inputs["input_ids"] = torch.cat([inputs["input_ids"], torch.tensor([tokenizer.sep_token_id])], dim=-1)
+    inputs["attention_mask"] = torch.cat([inputs["attention_mask"], torch.tensor([1])], dim=-1)
+    inputs = tokenizer.encode_plus(input_text, return_tensors="pt", add_special_tokens=False)
+    inputs["input_ids"] = torch.cat([inputs["input_ids"], inputs["input_ids"]], dim=-1)
+    inputs["attention_mask"] = torch.cat([inputs["attention_mask"], inputs["attention_mask"]], dim=-1)
 
     # Generate the response from the model with additional parameters
-    outputs = model.generate(**inputs, max_length=max_length, do_sample=True ,temperature=temperature)
+    outputs = model.generate(**inputs, max_length=max_length, do_sample=True, temperature=temperature)
 
     # Decode the response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()