Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,10 +1,8 @@
 import os
-
-# Rediriger les caches vers un dossier accessible en écriture
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache/huggingface/transformers'
-os.environ['HF_HOME'] = '/tmp/.cache/huggingface'
-os.environ['XDG_CACHE_HOME'] = '/tmp/.cache'
-os.environ['STREAMLIT_HOME'] = '/tmp/.streamlit'
+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
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -14,12 +12,12 @@ import torch
 def load_model():
     model_name = "openchat/openchat-3.5-0106"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-st.title("OpenChat 3.5 Demo")
+st.title("OpenChat - Demo")
 
 if "messages" not in st.session_state:
     st.session_state.messages = [{"role": "assistant", "content": "Posez-moi une question !"}]
@@ -28,20 +26,17 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=150)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-query = st.chat_input("Votre question ici...")
+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)
 
-    response = generate_response(query)
+    inputs = tokenizer(query, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=150)
+    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)