Analysis notebook

tasks/text.py

@@ -1,15 +1,16 @@
-from fastapi import APIRouter, Query
+import random
 from datetime import datetime
-from datasets import load_dataset
-from sklearn.metrics import accuracy_score
+
 import numpy as np
-import random
 import torch
-from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+from fastapi import APIRouter, Query
+from sklearn.metrics import accuracy_score
+from torch.utils.data import DataLoader, Dataset
 from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
 
+from .utils.emissions import clean_emissions_data, get_space_info, tracker
 from .utils.evaluation import TextEvaluationRequest
-from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
@@ -66,6 +67,8 @@ def bert_model(test_dataset: dict, model_type: str):
 
     if torch.cuda.is_available():
         device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
     else:
         device = torch.device("cpu")
     print("Using device:", device)