Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,45 +1,43 @@
 import os
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-# 🔧 Fix permissions & paths (HF Spaces & Streamlit)
 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['STREAMLIT_HOME'] = '/tmp/.hf/streamlit'
-os.environ['STREAMLIT_CONFIG_FILE'] = '/tmp/.hf/streamlit/config.toml'
 
-# 🧠 Model name (change if needed)
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"  
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-@st.cache_resource(show_spinner="🔄 Loading model...")
+@st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float32, 
-        device_map="auto"
-    )
+    model_name = "openchat/openchat-3.5-0106"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-st.title("💬 Simple LLM Chatbot (Streamlit + HF Transformers)")
-
-user_input = st.text_area("🧑‍💻 Posez votre question :", "", height=100)
-
-if st.button("Envoyer") and user_input.strip():
-    with st.spinner("✍️ Génération en cours..."):
-        input_ids = tokenizer.encode(user_input, return_tensors="pt")
-        output = model.generate(
-            input_ids,
-            max_new_tokens=150,
-            do_sample=True,
-            top_k=50,
-            top_p=0.95,
-            temperature=0.7
-        )
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-        st.markdown(f"**Réponse :**\n\n{response}")
+st.title("OpenChat 🤖")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = [{"role": "assistant", "content": "Salut ! Pose-moi une question."}]
+
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
+
+query = st.chat_input("Votre message...")
+
+if query:
+    st.session_state.messages.append({"role": "user", "content": query})
+    with st.chat_message("user"):
+        st.markdown(query)
+
+    inputs = tokenizer(query, return_tensors="pt").to(model.device)
+    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})
+    with st.chat_message("assistant"):
+        st.markdown(response)