Update app.py

app.py

@@ -7,21 +7,32 @@ import torch.nn as nn
 import re
 model_path = r'ssocean/NAIP'  
 
-model = AutoModelForSequenceClassification.from_pretrained(model_path, num_labels=1, load_in_8bit=True)
-
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+@spaces.GPU
+def init_model():
+    global model, tokenizer
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_path,
+        num_labels=1,
+        load_in_8bit=True,
+        device_map="auto"
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model.eval()
+    return model, tokenizer
 
+model, tokenizer = init_model()
 
 @spaces.GPU
 def predict(title, abstract):
-    global model
-    torch.cuda.set_device(torch.device('cuda:0'))
-    model.eval()
+    global model, tokenizer
+    
+    model_device = next(model.parameters()).device
     text = f'''Given a certain paper, Title: {title}\n Abstract: {abstract}. \n Predict its normalized academic impact (between 0 and 1):'''
-    inputs = tokenizer(text, return_tensors="pt")
-    inputs = inputs.to(torch.device('cuda:0'))
+    inputs = tokenizer(text, return_tensors="pt").to(model_device)
+
     with torch.no_grad():
         outputs = model(**inputs)
+
     probability = torch.sigmoid(outputs.logits).item()
     # reason for +0.05: We observed that the predicted values in the web demo are generally around 0.05 lower than those in the local deployment (due to differences in software/hardware environments). Therefore, we applied the following compensation in the web demo. Please do not use this in the local deployment.
     if probability + 0.05 >=1.0: