chore: import flash-attn first

app.py

@@ -1,3 +1,14 @@
+try:
+    import flash_attn
+except:
+    import subprocess
+
+    subprocess.run(
+        "pip install flash-attn --no-build-isolation",
+        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+        shell=True,
+    )
+
 import torch
 from transformers import (
     AutoModelForCausalLM,
@@ -19,18 +30,6 @@ except:
             return lambda x: x
 
 
-try:
-    import flash_attn
-except:
-    import subprocess
-
-    subprocess.run(
-        "pip install flash-attn --no-build-isolation",
-        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-        shell=True,
-    )
-
-
 MODEL_NAME = "hatakeyama-llm-team/Tanuki-8B-Instruct"
 
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)