Update app.py

app.py

@@ -15,6 +15,7 @@ from tqdm.auto import tqdm
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 
+from transformers import LlamaForConditionalGeneration
 
 
 
@@ -56,20 +57,19 @@ st.set_page_config(
 
 
 
-
 @st.cache_resource
 def load_llm_model():
-    # llm = HuggingFacePipeline.from_model_id(model_id= 'PyaeSoneK/LlamaV2LegalFineTuned', 
+    # llm = LlamaForConditionalGeneration.from_pretrained(model_id= 'PyaeSoneK/LlamaV2LegalFineTuned', 
     #                                         task= 'text2text-generation',
     #                                         model_kwargs={ "device_map": "auto",
     #                                                     "load_in_8bit": True,"max_length": 256, "temperature": 0,
     #                                                     "repetition_penalty": 1.5})
     
     
-    llm = HuggingFacePipeline.from_model_id(model_id= 'PyaeSoneK/LlamaV2LegalFineTuned', 
+    llm =  LlamaForConditionalGeneration.from_pretrained(model_id='PyaeSoneK/LlamaV2LegalFineTuned', 
                                         task= 'text2text-generation',
                                         
-                                        model_kwargs={ "max_length": 256, "temperature": 0,
+                                        model_kwargs={ "max_length": 128, "temperature": 0,
                                                       "torch_dtype":torch.float32,
                                                     "repetition_penalty": 1.3})
     return llm