Update app.py

app.py

@@ -4,11 +4,6 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
-# Load the model and tokenizer
-model_name = "joermd/llma-speedy"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
 # Random dog images for error messages
 random_dog = [
     "0f476473-2d8b-415e-b944-483768418a95.jpg",
@@ -16,21 +11,18 @@ random_dog = [
     # Add more images as needed
 ]
 
+# Function to reset conversation
 def reset_conversation():
     '''Resets conversation'''
     st.session_state.conversation = []
     st.session_state.messages = []
     return None
 
-# Create sidebar controls
+# Sidebar controls
 temp_values = st.sidebar.slider('Select a temperature value', 0.0, 1.0, 0.5)
 max_token_value = st.sidebar.slider('Select a max_token value', 1000, 9000, 5000)
 st.sidebar.button('Reset Chat', on_click=reset_conversation)
 
-# Set the model and display its name
-st.sidebar.write(f"You're now chatting with **{model_name}**")
-st.sidebar.markdown("*Generated content may be inaccurate or false.*")
-
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -40,28 +32,43 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# Accept user input
+# Set cache directory path
+cache_dir = "speedy-llm/"  # استبدل هذا بالمسار المخصص للتخزين في مساحة Hugging Face
+
+# Load model and tokenizer on-demand to save memory
 if prompt := st.chat_input(f"مرحبا انا سبيدي , كيف استطيع مساعدتك ؟"):
     with st.chat_message("user"):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
+    # Load model only when user submits a prompt
+    try:
+        # Load the tokenizer and model with caching in the specified directory
+        tokenizer = AutoTokenizer.from_pretrained("joermd/llma-speedy", cache_dir=cache_dir)
+        model = AutoModelForCausalLM.from_pretrained("joermd/llma-speedy", cache_dir=cache_dir)
+        
+        # Generate response
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(
+            inputs.input_ids,
+            max_new_tokens=max_token_value,
+            temperature=temp_values,
+            do_sample=True
+        )
+        assistant_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Clear memory (for CUDA) and delete the model to free up RAM
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        del model
+    
+    except Exception as e:
+        assistant_response = "😵‍💫 Connection issue! Try again later. Here's a 🐶:"
+        st.image(f'https://random.dog/{random_dog[np.random.randint(len(random_dog))]}')
+        st.write("Error message:")
+        st.write(e)
+
     # Display assistant response
     with st.chat_message("assistant"):
-        try:
-            inputs = tokenizer(prompt, return_tensors="pt")
-            outputs = model.generate(
-                inputs.input_ids,
-                max_new_tokens=max_token_value,
-                temperature=temp_values,
-                do_sample=True
-            )
-            assistant_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        except Exception as e:
-            assistant_response = "😵‍💫 Connection issue! Try again later. Here's a 🐶:"
-            st.image(f'https://random.dog/{random_dog[np.random.randint(len(random_dog))]}')
-            st.write("Error message:")
-            st.write(e)
-
-    st.markdown(assistant_response)
+        st.markdown(assistant_response)
     st.session_state.messages.append({"role": "assistant", "content": assistant_response})