Update app.py

app.py

@@ -49,9 +49,10 @@ with st.container():
         st.session_state.model = model
 
     #renders chat history
-    for message in st.session_state.chat_history: 
-        with st.chat_message(message["role"]):
-            st.write(message["content"])
+    for message in st.session_state.chat_history:
+        if(message["role"]!= "system"):
+            with st.chat_message(message["role"]):
+                st.write(message["content"])
             
     
     #Set up input text field
@@ -66,7 +67,7 @@ with st.container():
         #first_answer = chat_response.split("Human")[0] #Because of Predict it prints the whole conversation.Here we seperate the first answer only.
         tokenized_chat = tokenizer.apply_chat_template(st.session_state.chat_history, tokenize=True, add_generation_prompt=True, return_tensors="pt")
         #st.write(tokenizer.decode(tokenized_chat[0]))
-        outputs = model.generate(tokenized_chat.to('cuda'), max_new_tokens=128) 
+        outputs = model.generate(tokenized_chat, max_new_tokens=128) 
         first_answer = tokenizer.decode(outputs[0][tokenized_chat.shape[1]:],skip_special_tokens=True)
         
         with st.chat_message("assistant"):