Update app.py

app.py

@@ -15,6 +15,7 @@ DEFAULT_TEMPERATURE = 0.6
 DEFAULT_CHUNK_SIZE = 500
 DEFAULT_CHUNK_OVERLAP = 150
 DEFAULT_MAX_NEW_TOKENS = 256
+DEFAULT_OFFLOAD_DIR = "/tmp"
 
 # Set up logging
 logging.basicConfig(
@@ -23,7 +24,7 @@ logging.basicConfig(
 
 # Initialize and load the RAG model
 @st.cache_resource(show_spinner=False)
-def load_model(chat_model_id, embed_model_id, text_path, k, top_k, top_p, temperature, chunk_size, chunk_overlap, hf_token, max_new_tokens, quantization):
+def load_model(chat_model_id, embed_model_id, text_path, k, top_k, top_p, temperature, chunk_size, chunk_overlap, hf_token, max_new_tokens, quantization, offload_dir):
     rag_chain = BanglaRAGChain()
     rag_chain.load(
         chat_model_id=chat_model_id,
@@ -38,6 +39,7 @@ def load_model(chat_model_id, embed_model_id, text_path, k, top_k, top_p, temper
         hf_token=hf_token,
         max_new_tokens=max_new_tokens,
         quantization=quantization,
+        offload_dir=offload_dir,  # Pass the offload_dir here
     )
     return rag_chain
 
@@ -59,7 +61,7 @@ def main():
     text_path = st.sidebar.text_input("Text File Path", "text.txt")
     quantization = st.sidebar.checkbox("Enable Quantization (4-bit)", value=False)
     show_context = st.sidebar.checkbox("Show Retrieved Context", value=False)
-    offload_dir = st.sidebar.text_input("Offload Directory", "/tmp")  # Default to /tmp
+    offload_dir = st.sidebar.text_input("Offload Directory", DEFAULT_OFFLOAD_DIR)  # Default to /tmp
 
     # Load the model with the above configuration
     rag_chain = load_model(
@@ -75,7 +77,7 @@ def main():
         hf_token=None,  # If you're not using HF API token, set it to None
         max_new_tokens=max_new_tokens,
         quantization=quantization,
-        offload_dir=offload_dir,  # Pass offload dir here
+        offload_dir=offload_dir,  # Pass the offload_dir here
     )
     
     st.write("### Enter your question:")