Upload app.py with huggingface_hub

app.py

@@ -4,7 +4,6 @@ import logging
 import re
 
 import gradio as gr
-import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -29,12 +28,13 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
+    torch_dtype=dtype,
     trust_remote_code=True,
-)
+).to(device)
 model.eval()
 logger.info("Model loaded.")
 
@@ -73,7 +73,6 @@ def _grade_relevance(question: str, sources: list[dict]) -> bool:
     return top_score >= 0.02 or overlap >= 0.35
 
 
-@spaces.GPU
 def crag_answer(message: str, history: list[dict]) -> str:
     question = message.strip()
     if not question: