Update modules/pmbl.py

modules/pmbl.py

@@ -102,7 +102,7 @@ class PMBL:
             yield chunk
 
     def generate_response_task(self, system_prompt, prompt, n_ctx):
-        llm = Llama(model_path=self.model_path, n_ctx=n_ctx, n_threads=8, n_gpu_layers=-1, mlock=True)
+        llm = Llama(model_path=self.model_path, n_ctx=n_ctx, n_threads=8, n_gpu_layers=-1, use_flash_attn=True, mlock=True)
 
         response = llm(
             system_prompt,
@@ -110,9 +110,7 @@ class PMBL:
             temperature=0.7,
             stop=["</s>", "\nUser:", "\nuser:", "\nSystem:", "\nsystem:"],
             echo=False,
-            stream=True,
-            use_flash_attn=True,
-            mlock=True,
+            stream=True
         )
 
         response_text = ""