Update app.py

app.py

@@ -9,11 +9,32 @@ from pathlib import Path
 
 from transformers import pipeline, set_seed
 from sentence_transformers import SentenceTransformer
+import transformers.utils.import_utils as transformers_import_utils
 from datasets import load_dataset
 
+# Transformers 5 may query the CUDA stream while preparing CPU attention
+# masks (used internally when we call embedding_model.encode(...) on CPU).
+# In the GPU-less ZeroGPU parent process, that stray CUDA runtime call
+# leaves CUDA in a state that then breaks the next forked @spaces.GPU
+# worker's attempt to initialize its dynamically assigned GPU, causing
+# "RuntimeError: No CUDA GPUs are available" even though the code never
+# explicitly touches CUDA outside the decorated function.
+# Fix: skip only that query in the ZeroGPU parent; real GPU workers keep
+# the original Transformers behavior once CUDA_VISIBLE_DEVICES is set.
+# See: https://huggingface.co/spaces/zero-gpu-explorers/README/discussions/181
+if hasattr(transformers_import_utils, "is_cuda_stream_capturing"):
+    _original_is_cuda_stream_capturing = transformers_import_utils.is_cuda_stream_capturing
+
+    def _zerogpu_safe_is_cuda_stream_capturing(*args, **kwargs):
+        if os.getenv("SPACES_ZERO_GPU") == "1" and not os.getenv("CUDA_VISIBLE_DEVICES"):
+            return False
+        return _original_is_cuda_stream_capturing(*args, **kwargs)
+
+    transformers_import_utils.is_cuda_stream_capturing = _zerogpu_safe_is_cuda_stream_capturing
+
 # 1. Initialize LLM
 print("Loading BioMistral LLM...")
-llm = None
+llm = pipeline('text-generation', model='BioMistral/BioMistral-7B', torch_dtype=torch.bfloat16, device_map="auto")
 
 # 2. Load Datasets
 print("Loading datasets...")
@@ -78,15 +99,6 @@ Question:
 @spaces.GPU(duration=60)
 def generate_answer(prompt: str, max_new_tokens: int = 80) -> str:
     set_seed(42)
-    global llm
-    if llm is None:
-        print("Initializing BioMistral on ZeroGPU...")
-        llm = pipeline(
-            'text-generation',
-            model='BioMistral/BioMistral-7B',
-            torch_dtype=torch.bfloat16,
-            device_map="auto"
-        )
     if llm.model.config.pad_token_id is None:
         llm.model.config.pad_token_id = llm.model.config.eos_token_id