Update app.py

app.py

@@ -162,15 +162,14 @@ def _load_embed_and_rerank_on_gpu():
 def _load_generation_model_on_gpu():
     """
     Try Nemotron 30B FP8 first.
-    If it fails (e.g., missing mamba-ssm), fall back to a smaller text model.
+    If it fails, fall back to a smaller text model.
     """
     global _gen_model, _gen_tokenizer
-
     if _gen_model is not None and _gen_tokenizer is not None:
         return _gen_model, _gen_tokenizer
 
     device = _cuda()
-
+    
     # 1) Try Nemotron FP8
     try:
         print("[INFO] Lazy-loading GENERATION model (Nemotron 30B FP8) on GPU...")
@@ -179,33 +178,41 @@ def _load_generation_model_on_gpu():
             trust_remote_code=True,
             use_fast=True,
         )
+        
+        # FIX: We force 'eager' or 'flash_attention_2' because this model 
+        # doesn't support the default 'sdpa' implementation yet.
+        # Since you installed the FA2 wheels, we'll try to use that first.
+        gen_attn_impl = ATTN_IMPL if ATTN_IMPL == "flash_attention_2" else "eager"
+
         _gen_model = AutoModelForCausalLM.from_pretrained(
             GENERATION_MODEL_ID,
             trust_remote_code=True,
             torch_dtype="auto",
             device_map="auto",
-            attn_implementation=ATTN_IMPL,
+            attn_implementation=gen_attn_impl, # Changed here
         ).eval()
-        print("[INFO] Nemotron generation model loaded OK")
+        
+        print(f"[INFO] Nemotron generation model loaded OK with {gen_attn_impl}")
         return _gen_model, _gen_tokenizer
 
     except Exception as e:
         print(f"[WARN] Nemotron FP8 load failed: {repr(e)}")
         print(f"[WARN] Falling back to: {FALLBACK_GEN_MODEL_ID}")
-
+        
         _gen_tokenizer = AutoTokenizer.from_pretrained(
             FALLBACK_GEN_MODEL_ID,
             trust_remote_code=True,
             use_fast=True,
         )
+        
         _gen_model = AutoModelForCausalLM.from_pretrained(
             FALLBACK_GEN_MODEL_ID,
             trust_remote_code=True,
             torch_dtype=torch.bfloat16,
             device_map="auto",
-            attn_implementation=ATTN_IMPL,
+            attn_implementation="sdpa", # Fallback model usually supports SDPA
         ).eval()
-
+        
         return _gen_model, _gen_tokenizer