Update app.py

app.py

@@ -7,10 +7,14 @@ import torch
 import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+
+offload_folder = '/path/to/your/offload/folder'
+model = YourModel(..., offload_folder=offload_folder)
 model = AutoModelForCausalLM.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
                                     device_map='auto',
                                   torch_dtype=torch.float16,
-                                  use_auth_token= st.secrets['hf_access_token'],)
+                                  use_auth_token= st.secrets['hf_access_token'],
+                                    offload_folder=offload_folder,)
                                             #  load_in_4bit=True
 
 tokenizer = AutoTokenizer.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",