Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,8 @@
 import os
+os.environ['HOME'] = '/tmp'
+os.environ['XDG_CACHE_HOME'] = '/tmp/.cache'
 os.environ['HF_HOME'] = '/tmp/.hf'
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/.hf/transformers'
-os.environ['XDG_CACHE_HOME'] = '/tmp/.hf/cache'
 os.environ['STREAMLIT_HOME'] = '/tmp/.hf/streamlit'
 
 import streamlit as st
@@ -17,14 +18,14 @@ def load_model():
 
 tokenizer, model = load_model()
 
-st.title("OpenChat - Demo")
+st.title("OpenChat 🤖")
 
 if "messages" not in st.session_state:
-    st.session_state.messages = [{"role": "assistant", "content": "Posez-moi une question !"}]
+    st.session_state.messages = [{"role": "assistant", "content": "Salut ! Pose-moi une question."}]
 
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
 
 query = st.chat_input("Votre message...")
 
@@ -34,7 +35,7 @@ if query:
         st.markdown(query)
 
     inputs = tokenizer(query, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=150)
+    outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True, top_p=0.95, top_k=50)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     st.session_state.messages.append({"role": "assistant", "content": response})