Upload main.py

main.py

@@ -4,16 +4,16 @@ import uvicorn
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 app = FastAPI()
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-n_gpu = torch.cuda.device_count()
-torch.cuda.get_device_name(0)
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# n_gpu = torch.cuda.device_count()
+# torch.cuda.get_device_name(0)
 tokenizer = AutoTokenizer.from_pretrained("SuiGio/roberta_pubmesh")
 model = AutoModelForSequenceClassification.from_pretrained("SuiGio/roberta_pubmesh")
 
-try:
-    model.cuda()
-except:
-    pass
+# try:
+#     model.cuda()
+# except:
+#     pass
 
 
 label_names= [ "Anatomy [A]",
@@ -73,8 +73,8 @@ user_text = " BACKGROUND: The distal GI microbiota of hospitalized preterm neona
 
 def get_preds(user_text, model=model, tokenizer=tokenizer, label_names=label_names):
     inputs = tokenizer(user_text, padding=True, truncation=True, max_length=128, return_tensors="pt")
-    input_ids = inputs["input_ids"].to(device)
-    attention_mask = inputs["attention_mask"].to(device)
+    input_ids = inputs["input_ids"]#.to(device)
+    attention_mask = inputs["attention_mask"]#.to(device)
     model.eval()
 
     # Forward pass with the model