Update app.py

app.py

@@ -5,6 +5,7 @@ A math reading coach for grades 3–8.
 
 import os
 
+import spaces
 import gradio as gr
 import torch
 from peft import PeftModel
@@ -50,24 +51,14 @@ EXAMPLES = [
 
 print(f"Loading base model: {BASE_MODEL_ID}")
 
-if torch.cuda.is_available():
-    model_dtype = torch.float16
-    device_map = "auto"
-    print("GPU detected. Loading model in float16.")
-else:
-    model_dtype = torch.float32
-    device_map = None
-    print("No GPU detected. CPU loading may be slow or require more memory.")
-
-
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL_ID,
-    torch_dtype=model_dtype,
-    device_map=device_map,
+    torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     **HUB_ARGS,
 )
 
+base_model = base_model.to("cuda")
 
 print(f"Loading Estella adapter: {ADAPTER_ID}")
 
@@ -77,21 +68,26 @@ model = PeftModel.from_pretrained(
     **HUB_ARGS,
 )
 
+model = model.to("cuda")
+model.eval()
+
 tokenizer = AutoTokenizer.from_pretrained(
     ADAPTER_ID,
     **HUB_ARGS,
 )
 
-model.eval()
-
 print("Estella is ready.")
 
 
+@spaces.GPU(duration=60)
 def ask_estella(message: str, history: list) -> str:
     """Return one reading hint without solving the math problem."""
 
     if not message or not message.strip():
-        return "Please type a math word problem and I’ll help explain what it is asking!"
+        return (
+            "Please type a math word problem and I’ll help explain "
+            "what it is asking!"
+        )
 
     messages = [
         {
@@ -120,6 +116,7 @@ def ask_estella(message: str, history: list) -> str:
     )
 
     input_device = next(model.parameters()).device
+
     inputs = {
         key: value.to(input_device)
         for key, value in inputs.items()