Update app.py

app.py

@@ -66,8 +66,8 @@ 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, max_new_tokens=128) 
-        first_answer = tokenizer.decode(outputs[0],skip_special_tokens=True)
+        outputs = model.generate(tokenized_chat.to('cuda'), max_new_tokens=128) 
+        first_answer = tokenizer.decode(outputs[0][tokenized_chat.shape[1]:],skip_special_tokens=True)
         
         with st.chat_message("assistant"):
             st.write(first_answer)