Update app.py

app.py

@@ -80,6 +80,15 @@ def load_llm_model():
     #                                              use_auth_token= st.secrets['hf_access_token'],
                                                   
     #return llm
+    model = 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,)
+
     pipe = pipeline("text-generation",
                 model=model,
                 tokenizer= tokenizer,
@@ -92,18 +101,7 @@ def load_llm_model():
                 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
+    return model