Update app.py

app.py

@@ -1,5 +1,8 @@
 import os
 import subprocess
+
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
 import threading
 
 # subprocess.check_call([os.sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
@@ -22,6 +25,7 @@ def load_model():
         MODEL_ID,
         torch_dtype=torch.bfloat16,  # if device == "cuda" else torch.float32,
         device_map="auto",  # if device == "cuda" else None,
+        attn_implementation="flash_attention_2",
     )  # .cuda()
     print(f"Selected device:", device)
     return model, tokenizer, device