Allow user to choose device for models

outlines/models/transformers.py

@@ -8,7 +8,6 @@
 if TYPE_CHECKING:
     from transformers import PreTrainedModel, PreTrainedTokenizer
 
-
 __all__ = ["transformers"]
 
 
@@ -19,10 +18,9 @@ def __init__(
         self,
         model: "PreTrainedModel",
         tokenizer: "PreTrainedTokenizer",
-        device: Optional[str] = None,
     ):
-        self.device = device if device is not None else "cpu"
-        self.model = model.to(self.device)
+        self.device = model.device
+        self.model = model
         self.tokenizer = tokenizer
 
     def __call__(
@@ -86,15 +84,42 @@ def convert_token_to_string(self, token: str) -> str:
         return string
 
 
-def transformers(model_name: str, device: Optional[str] = None, **model_kwargs):
+def transformers(
+    model_name: str,
+    device: Optional[str] = None,
+    model_kwargs: dict = {},
+    tokenizer_kwargs: dict = {},
+):
+    """Instantiate a model from the `transformers` library and its tokenizer.
+
+    Parameters
+    ----------
+    model_name
+        The name of the model as listed on Hugging Face's model page.
+    device_map
+        The device(s) on which the model should be loaded. This overrides
+        the value passed for `device_map` in `model_kwargs`.
+    model_kwargs
+        A dictionary that contains the keyword arguments to pass to the
+        `from_pretrained` method when loading the model.
+    tokenizer_kwargs
+        A dictionary that contains the keyword arguments to pass to the
+        `from_pretrained` method when loading the tokenizer.
+
+    Returns
+    -------
+    A `TransformersModel` model instance.
+
+    """
     try:
         from transformers import AutoModelForCausalLM
     except ImportError:
         raise ImportError(
             "The `transformers` library needs to be installed in order to use `transformers` models."
         )
 
+    model_kwargs["device_map"] = device
     model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
-    tokenizer = TransformersTokenizer(model_name)
+    tokenizer = TransformersTokenizer(model_name, **tokenizer_kwargs)
 
-    return Transformers(model, tokenizer, device)
+    return Transformers(model, tokenizer)

tests/models/test_transformers.py

@@ -39,12 +39,12 @@ def test_tokenizer():
 
 
 def test_model():
-    with pytest.raises(RuntimeError, match="Expected one of cpu, cuda"):
+    with pytest.raises(ValueError, match="When passing device_map as a string"):
         transformers(TEST_MODEL, device="non_existent")
 
     model = transformers(TEST_MODEL, device="cpu")
     assert isinstance(model.tokenizer, TransformersTokenizer)
-    assert model.device == "cpu"
+    assert model.device.type == "cpu"
 
     input_ids = torch.tensor([[0, 1, 2]])
     logits = model(input_ids, torch.ones_like(input_ids))