CrossEncoder device

sentence_transformers/SentenceTransformer.py

@@ -4,7 +4,7 @@
 import shutil
 from collections import OrderedDict
 import warnings
-from typing import List, Dict, Tuple, Iterable, Type, Union, Callable, Optional, Literal, TYPE_CHECKING
+from typing import List, Dict, Tuple, Iterable, Type, Union, Callable, Optional, TYPE_CHECKING
 import numpy as np
 from numpy import ndarray
 import transformers
@@ -29,6 +29,7 @@
     load_dir_path,
     load_file_path,
     save_to_hub_args_decorator,
+    get_device_name,
 )
 from .models import Transformer, Pooling, Normalize
 from .model_card_templates import ModelCardTemplate
@@ -41,23 +42,6 @@
     from sentence_transformers.readers import InputExample
 
 
-def get_device_name() -> Literal["mps", "cuda", "cpu"]:
-    """
-    Returns the name of the device where this module is running on.
-    It's simple implementation that doesn't cover cases when more powerful GPUs are available and
-    not a primary device ('cuda:0') or MPS device is available, but not configured properly:
-    https://pytorch.org/docs/master/notes/mps.html
-
-    :return: Device name, like 'cuda' or 'cpu'
-    """
-    if torch.cuda.is_available():
-        return "cuda"
-    elif torch.backends.mps.is_available():
-        return "mps"
-    else:
-        return "cpu"
-
-
 class SentenceTransformer(nn.Sequential):
     """
     Loads or creates a SentenceTransformer model that can be used to map sentences / text to embeddings.

sentence_transformers/cross_encoder/CrossEncoder.py

@@ -10,6 +10,7 @@
 from tqdm.autonotebook import tqdm, trange
 from .. import SentenceTransformer, util
 from ..evaluation import SentenceEvaluator
+from ..util import get_device_name
 
 
 logger = logging.getLogger(__name__)
@@ -73,7 +74,7 @@ def __init__(
         self.max_length = max_length
 
         if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
+            device = get_device_name()
             logger.info("Use pytorch device: {}".format(device))
 
         self._target_device = torch.device(device)

sentence_transformers/util.py

@@ -1,7 +1,7 @@
 import functools
 import requests
 from torch import Tensor, device
-from typing import List, Callable
+from typing import List, Callable, Literal
 from tqdm.autonotebook import tqdm
 import sys
 import importlib
@@ -556,3 +556,20 @@ def wrapper(self, *args, **kwargs):
         return func(self, *args, **kwargs)
 
     return wrapper
+
+
+def get_device_name() -> Literal["mps", "cuda", "cpu"]:
+    """
+    Returns the name of the device where this module is running on.
+    It's simple implementation that doesn't cover cases when more powerful GPUs are available and
+    not a primary device ('cuda:0') or MPS device is available, but not configured properly:
+    https://pytorch.org/docs/master/notes/mps.html
+
+    :return: Device name, like 'cuda' or 'cpu'
+    """
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    else:
+        return "cpu"