Update app.py

app.py

@@ -6,20 +6,21 @@ from transformers import AutoModel
 import torch
 import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 # App framework
 st.title('🦜Seon\'s Legal QA For Dummies 🔗 ')
 
-
-
-offload_folder = 'C:\model_weights'
-
-model = AutoModelForCausalLM.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
+with init_empty_weights():
+    model = AutoModelForCausalLM.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
                                     device_map='auto',
                                   torch_dtype=torch.float16,
                                   use_auth_token= st.secrets['hf_access_token'],
-                                    offload_folder=offload_folder,
                                     )
+
+model = load_checkpoint_and_dispatch(
+    model, checkpoint=checkpoint_file, device_map="auto"
+)
                                             #  load_in_4bit=True
 
 tokenizer = AutoTokenizer.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",