Set device for local PyTorchGemmaLanguageModel

concordia/language_model/pytorch_gemma_model.py

@@ -35,6 +35,7 @@ def __init__(
       *,
       measurements: measurements_lib.Measurements | None = None,
       channel: str = language_model.DEFAULT_STATS_CHANNEL,
+      device: str = 'cpu'
   ) -> None:
     """Initializes the instance.
 
@@ -46,11 +47,12 @@ def __init__(
     """
     self._model_name = model_name
     self._tokenizer_name = model_name
+    self._device = device
 
     os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 
     self._model = transformers.GemmaForCausalLM.from_pretrained(
-        self._model_name)
+        self._model_name).to(self._device)
     self._tokenizer = transformers.AutoTokenizer.from_pretrained(
         self._tokenizer_name)
 
@@ -80,14 +82,14 @@ def sample_text(
     inputs = self._tokenizer(prompt_with_system_message, return_tensors='pt')
 
     generated_tokens = self._model.generate(
-        inputs.input_ids,
+        inputs.input_ids.to(self._device),
         max_new_tokens=max_tokens,
         return_dict_in_generate=True,
         output_scores=True,
     )
 
     response = self._tokenizer.decode(
-        np.int64(generated_tokens.sequences[0]),
+        np.int64(generated_tokens.sequences[0].cpu()),
         skip_special_tokens=True,
         clean_up_tokenization_spaces=False,
     )
@@ -116,13 +118,13 @@ def sample_choice(
 
     inputs = self._tokenizer(prompt, return_tensors='pt')
     generated_tokens = self._model.generate(
-        inputs.input_ids,
+        inputs.input_ids.to(self._device),
         max_new_tokens=1,
         return_dict_in_generate=True,
         output_scores=True,
     )
     sample = self._tokenizer.batch_decode(
-        [np.argmax(generated_tokens.scores[0][0])],
+        [np.argmax(generated_tokens.scores[0][0].cpu())],
         skip_special_tokens=True,
         clean_up_tokenization_spaces=False)[0]
     answer = sampling.extract_choice_response(sample)

concordia/language_model/utils.py

@@ -32,6 +32,7 @@ def language_model_setup(
     model_name: str,
     api_key: str | None = None,
     disable_language_model: bool = False,
+    device: str = 'cpu'
 ) -> language_model.LanguageModel:
   """Get the wrapped language model.
 
@@ -69,7 +70,7 @@ def language_model_setup(
   elif api_type == 'openai':
     return gpt_model.GptLanguageModel(model_name, api_key=api_key)
   elif api_type == 'pytorch_gemma':
-    return pytorch_gemma_model.PyTorchGemmaLanguageModel(model_name)
+    return pytorch_gemma_model.PyTorchGemmaLanguageModel(model_name, device=device)
   elif api_type == 'together_ai':
     return together_ai.Gemma2(model_name, api_key=api_key)
   else: