Update app.py

app.py

@@ -15,7 +15,10 @@ from tqdm.auto import tqdm
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 
-from transformers import LlamaForConditionalGeneration
+from transformers import AutoModel
+import torch
+import transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 
@@ -66,11 +69,12 @@ def load_llm_model():
     #                                                     "repetition_penalty": 1.5})
     
     
-    llm =  LlamaForConditionalGeneration.from_pretrained(model_id='PyaeSoneK/LlamaV2LegalFineTuned', 
+    llm =  AutoModelForCausalLM.from_pretrained('PyaeSoneK/LlamaV2LegalFineTuned', 
                                         task= 'text2text-generation',
                                         
                                         model_kwargs={ "max_length": 128, "temperature": 0,
                                                       "torch_dtype":torch.float32,
+                                                      use_auth_token=True,
                                                     "repetition_penalty": 1.3})
     return llm