Update app.py

app.py

@@ -1,3 +1,9 @@
+import subprocess
+subprocess.run(
+    'pip install flash-attn --no-build-isolation',
+    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+    shell=True
+)
 import os
 import time
 import spaces
@@ -36,7 +42,8 @@ h3 {
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID, 
-    torch_dtype=torch.float16, 
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2",
     trust_remote_code=True).cuda()
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)