Update app.py

app.py

@@ -57,7 +57,12 @@ st.set_page_config(
     page_icon = '🕵')
 
 
+from transformers import AutoModel
+import torch
+import transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
+from transformers import pipeline
 
 @st.cache_resource
 def load_llm_model():
@@ -67,13 +72,39 @@ def load_llm_model():
     #                                                     "load_in_8bit": True,"max_length": 256, "temperature": 0,
     #                                                     "repetition_penalty": 1.5})
     
-    token = st.secrets['hf_access_token']
-    llm =  AutoModelForCausalLM.from_pretrained('PyaeSoneK/LlamaV2LegalFineTuned', 
-                                                  device_map='auto',
-                                                  torch_dtype=torch.float16,
-                                                  use_auth_token= st.secrets['hf_access_token'],
+    #token = st.secrets['hf_access_token']
+    #llm =  AutoModelForCausalLM.from_pretrained(model_id = 'PyaeSoneK/LlamaV2LegalFineTuned', 
+    #                                           task = 'text2text-generation',
+    #                                                device_map='auto',
+    #                                              torch_dtype=torch.float16,
+    #                                              use_auth_token= st.secrets['hf_access_token'],
                                                   )
-    return llm
+    #return llm
+    pipe = pipeline("text-generation",
+                model=model,
+                tokenizer= tokenizer,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                max_new_tokens = 512,
+                do_sample=True,
+                top_k=30,
+                num_return_sequences=1,
+                eos_token_id=tokenizer.eos_token_id
+                )
+
+
+
+    llm = AutoModelForCausalLM.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
+                                   device_map='auto',
+                                  torch_dtype=torch.float16,
+                                  use_auth_token= st.secrets['hf_access_token'],)
+                                            #  load_in_4bit=True
+
+    tokenizer = AutoTokenizer.from_pretrained("PyaeSoneK/LlamaV2LegalFineTuned",
+                                          use_auth_token=True,)
+
+return llm
+
 
 
 @st.cache_resource