Update app.py

app.py

@@ -11,19 +11,18 @@ from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 # App framework
 st.title('🦜Seon\'s Legal QA For Dummies 🔗 ')
 
-with init_empty_weights():
-    model = AutoModelForCausalLM.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
+
+model = AutoModelForCausalLM.from_pretrained("PyaeSoneK/pythia_70m_legalQA",
                                     device_map='auto',
                                   torch_dtype=torch.float16,
                                   use_auth_token= st.secrets['hf_access_token'],
                                     )
 
-model = load_checkpoint_and_dispatch(
-    model, checkpoint=checkpoint_file, device_map="auto"
-)
+
+  
                                             #  load_in_4bit=True
 
-tokenizer = AutoTokenizer.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
+tokenizer = AutoTokenizer.from_pretrained("PyaeSoneK/pythia_70m_legalQA",
                                           use_auth_token=st.secrets['hf_access_token'],)
 
 # Use a pipeline for later
@@ -44,7 +43,7 @@ pipe = pipeline("text-generation",
 import json
 import textwrap
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+
 
 B_INST, E_INST = "[INST]", "[/INST]"
 B_SYS, E_SYS = "<>\n", "\n<>\n\n"