Update app.py

app.py

@@ -5,10 +5,8 @@ import subprocess
 import spaces
 
 
-@spaces.GPU
-def _build_flash_attn():
-    subprocess.check_call("pip install flash-attn", shell=True)
-_build_flash_attn() # This is how we'll build flash-attn.
+# Install flash-attn
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 # Initialize the model pipeline
 generator = pipeline('text-generation', model='mistralai/Mistral-7B-v0.1', torch_dtype=torch.bfloat16, use_flash_attention_2=True)
 @spaces.GPU