V1.1

app.py

@@ -2,9 +2,14 @@
 Hugging Face Spaces entry point.
 Runs the MedGemma triage Gradio demo. Set HF_TOKEN in Space secrets for gated model access.
 """
+import os
 import sys
 from pathlib import Path
 
+# Avoid "Invalid value for environment variable OMP_NUM_THREADS" on Spaces (set before any OpenMP use)
+if "OMP_NUM_THREADS" not in os.environ or not str(os.environ.get("OMP_NUM_THREADS", "")).strip().isdigit():
+    os.environ["OMP_NUM_THREADS"] = "1"
+
 # Ensure project root is on path (Spaces run from repo root)
 ROOT = Path(__file__).resolve().parent
 if str(ROOT) not in sys.path:

src/models/medgemma_client.py

@@ -147,12 +147,27 @@ class MedGemmaClient:
         # Set default parameters
         max_length = max_length or ModelConfig.MAX_LENGTH
         max_new_tokens = max_new_tokens or ModelConfig.MAX_NEW_TOKENS
-        temperature = temperature or ModelConfig.TEMPERATURE
+        temperature = temperature if temperature is not None else ModelConfig.TEMPERATURE
         top_p = top_p or ModelConfig.TOP_P
         top_k = top_k or ModelConfig.TOP_K
         
-        try:
-            # Tokenize input
+        # Avoid CUDA "probability tensor contains inf/nan or element < 0" by using greedy when temp is low
+        # and clamping temperature when sampling (very low temp can make logits explode in float16)
+        do_sample = float(temperature) > 0.05
+        if do_sample:
+            temperature = max(0.05, min(float(temperature), 2.0))
+        
+        def _run_generate(do_samp: bool, temp: float) -> str:
+            gen_config = GenerationConfig(
+                max_new_tokens=max_new_tokens,
+                temperature=temp if do_samp else 1.0,
+                top_p=top_p if do_samp else 1.0,
+                top_k=top_k if do_samp else 0,
+                do_sample=do_samp,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                **kwargs
+            )
             inputs = self.tokenizer(
                 prompt,
                 return_tensors="pt",
@@ -160,39 +175,30 @@ class MedGemmaClient:
                 truncation=True,
                 max_length=max_length
             ).to(self.device)
-            
-            # Create generation config
-            generation_config = GenerationConfig(
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                do_sample=temperature > 0,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
-                **kwargs
-            )
-            
-            # Generate
             with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    generation_config=generation_config
-                )
-            
-            # Decode output
-            generated_text = self.tokenizer.decode(
-                outputs[0],
-                skip_special_tokens=True
-            )
-            
-            # Remove the input prompt from output
+                outputs = self.model.generate(**inputs, generation_config=gen_config)
+            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             if generated_text.startswith(prompt):
                 generated_text = generated_text[len(prompt):].strip()
-            
             return generated_text
-            
+        
+        try:
+            return _run_generate(do_sample, temperature)
         except Exception as e:
+            err_msg = str(e).lower()
+            is_cuda_assert = (
+                "cuda" in err_msg
+                or "device-side assert" in err_msg
+                or "acceleratorerror" in err_msg
+                or "probability tensor" in err_msg
+            )
+            if is_cuda_assert and do_sample:
+                logger.warning(f"Generation failed with device error, retrying with greedy decoding: {e}")
+                try:
+                    return _run_generate(do_samp=False, temp=1.0)
+                except Exception as retry_e:
+                    logger.error(f"Generation failed after greedy fallback: {retry_e}")
+                    raise retry_e
             logger.error(f"Generation failed: {e}")
             raise