Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
 
 # Create a class for the session state
 class SessionState:
@@ -22,37 +21,32 @@ with st.spinner('Wait for it... the model is loading'):
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Create a chat input for the user
-input_text = st.chat_input("Enter your message:")
+# Create a text input for the user
+input_text = st.text_input("Enter your message:")
 
 # Check if the user has entered a message
 if input_text:
     # Add the user's message to the conversation history
-    session_state.conversation_history.append(("User", input_text))
+    session_state.conversation_history.append(input_text)
 
     # Create conversation history string
-    history_string = "\n".join(message for role, message in session_state.conversation_history)
+    history_string = "\n".join(session_state.conversation_history)
 
     # Tokenize the input text and history
-    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)
+    inputs = tokenizer.encode_plus(history_string, input_text, return_tensors="pt")
 
     # 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, temperature=temperature)
 
     # Decode the response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
     # Add the model's response to the conversation history
-    session_state.conversation_history.append(("Assistant", response))
-
-    # Display the conversation history using st.chat
-    for role, message in session_state.conversation_history:
-        if role == "User":
-            st.chat_message(message, is_user=True)
-        else:
-            st.chat_message(message, is_user=False)
+    session_state.conversation_history.append(response)
+
+    # Display the conversation history
+    st.write("Conversation History:")
+    for i in range(0, len(session_state.conversation_history), 2):
+        st.write("User: " + session_state.conversation_history[i])
+        if i+1 < len(session_state.conversation_history):
+            st.write("Assistant: " + session_state.conversation_history[i+1])