Update app.py

app.py

@@ -3,22 +3,27 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 import requests
 
+# Ensure torch is installed before flash_attn
+try:
+    import flash_attn
+except ImportError:
+    raise RuntimeError("Missing required dependency: flash_attn. Install with `pip install flash-attn`")
+
 # Load DeepSeek-R1 model with trust_remote_code enabled
 model_name = "deepseek-ai/DeepSeek-R1"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-# Ensure compatibility with `flash_attn` and force proper dtype
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,  # Forces float16 to prevent fp8 issue
-        device_map="auto",
-        trust_remote_code=True
-    )
-except ImportError as e:
-    raise RuntimeError("Missing required dependency: flash_attn. Install with `pip install flash_attn`") from e
-
-# Use a text-generation pipeline for better inference
+# Ensure compatibility and force execution on GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,  # Prevents fp8 errors
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# Use a text-generation pipeline
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)