faster inference :)

app.py

@@ -5,10 +5,17 @@ from transformers import StoppingCriteria, StoppingCriteriaList
 import torch
 import spaces
 import os
+import subprocess
 
+# Install flash attention
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 model_name = "microsoft/Phi-3-medium-128k-instruct"
 from transformers import AutoModelForCausalLM, AutoTokenizer
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map='cuda', torch_dtype=torch.float16, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map='cuda', _attn_implementation="flash_attention_2", torch_dtype=torch.float16, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 class StopOnTokens(StoppingCriteria):
@@ -19,7 +26,7 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 model.to('cuda')
-@spaces.GPU()
+@spaces.GPU(queue=False)
 def predict(message, history, temperature, max_tokens, top_p, top_k):
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()