Update app.py

app.py

@@ -4,20 +4,25 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import warnings
+from huggingface_hub import spaces
 
 # Suppress all warnings
 warnings.filterwarnings("ignore")
 
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
-# Handle GPU initialization safely
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Initialize GPU for Hugging Face Spaces
+@spaces.GPU
+def init_gpu():
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Initialize model and tokenizer
 MODEL_NAME = "s-nlp/roberta-base-formality-ranker"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
-# Move model to appropriate device
+# Move model to GPU
+device = init_gpu()
 model = model.to(device)
 
 app = FastAPI(title="Formality Classifier API")