Update app.py

app.py

@@ -47,12 +47,51 @@ RERANK_COMMIT_HASH = "47e5a355d1a050c3e5f69d53f14964b1d34bcd9d"
 GENERATION_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
 FALLBACK_GEN_MODEL_ID = os.getenv("FALLBACK_GEN_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
 
-ATTN_IMPL = "flash_attention_2" if os.getenv("USE_FA2", "0") == "1" else "sdpa"
+# ATTN_IMPL = "flash_attention_2" if os.getenv("USE_FA2", "0") == "1" else "sdpa"
 
 modality_to_tokens = {"image": 2048, "image_text": 10240, "text": 8192}
 
 PATH_TO_EMBEDDING_FILE = os.getenv("EMBEDDINGS_FILE", "image_text_embeddings_10k.safetensors")
 
+
+def check_flash_attention():
+    import torch
+    from transformers.utils import is_flash_attn_2_available
+
+    print(f"--- Flash Attention Check ---")
+    print(f"PyTorch version: {torch.__version__}")
+    print(f"CUDA available: {torch.cuda.is_available()}")
+    
+    # Transformers helper check
+    fa2_available = is_flash_attn_2_available()
+    print(f"Transformers reports FA2 available: {fa2_available}")
+
+    if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability()
+        print(f"GPU Compute Capability: {capability}")
+        if capability[0] < 8:
+            print("Note: FA2 requires Compute Capability 8.0+ (Ampere or newer).")
+    
+    return fa2_available
+
+# Determine best implementation
+if check_flash_attention():
+    ATTN_IMPL = "flash_attention_2"
+else:
+    ATTN_IMPL = "sdpa" # Fallback to Scaled Dot Product Attention
+
+print(f"[INFO] Using {ATTN_IMPL} for model loading.")
+
+# model = AutoModelForCausalLM.from_pretrained(
+#     model_id,
+#     torch_dtype=torch.float16, # FA2 requires fp16 or bf16
+#     attn_implementation=best_attn,
+#     trust_remote_code=True
+# ).to("cuda")
+
+# Call it inside your setup or first GPU call
+# FLASH_AVAILABLE = check_flash_attention()
+
 # -----------------------------------------------------------------------------
 # Load dataset + embeddings (CPU only)
 # -----------------------------------------------------------------------------