Update tasks/text.py

tasks/text.py

@@ -24,6 +24,12 @@ DESCRIPTION = "First Baseline"
 ROUTE = "/text"
 
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+
 MODEL = "mlp" #mlp, ct, modern
 
 class ConspiracyClassification(
@@ -125,7 +131,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     if MODEL =="mlp":
         model = ConspiracyClassification.from_pretrained("ypesk/frugal-ai-mlp-baseline")
-        
+        model = model.to(device)        
         emb_model = SentenceTransformer("paraphrase-MiniLM-L3-v2")
         batch_size = 6
 
@@ -136,6 +142,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         
     elif MODEL == "ct":
         model = CovidTwitterBertClassifier.from_pretrained("ypesk/ct-baseline")
+        model = model.to(device)
         tokenizer = AutoTokenizer.from_pretrained('digitalepidemiologylab/covid-twitter-bert')
         
         test_texts = [t['quote'] for t in test_dataset]
@@ -158,7 +165,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     model.eval()
     predictions = []
     for batch in tqdm(test_dataloader):
-        
+        batch = tuple(t.to(device) for t in batch)
         with torch.no_grad():
             if MODEL =="mlp":
                 b_texts = batch[0]