Support pygrain.MapDataset in dpsynth.text DPFineTuner.

dpsynth/text/dp_sft.py

@@ -47,8 +47,10 @@
 from dpsynth.text import model
 from gemma import gm
 from gemma import peft
+import grain.python as pygrain
 from jax_privacy import execution_plan
 from jax_privacy import training
+import numpy as np
 import optax
 
 
@@ -75,7 +77,7 @@ class DPFineTuner(api.DPMechanism):
   A ``DPMechanism`` that wraps ``DPTrainer`` with tokenization, model loading,
   and LoRA handling. All configuration is specified at construction time;
   ``__call__`` takes ``(rng, data)`` where ``data`` is a sequence of text
-  strings.
+  string pairs or a pygrain MapDataset.
   """
 
   model_variant: model.GemmaModel
@@ -134,23 +136,36 @@ def dp_event(self) -> dp_accounting.DpEvent:
   def __call__(
       self,
       rng: int,
-      data: Sequence[tuple[str, str]],
+      data: Sequence[tuple[str, str]] | pygrain.MapDataset,
   ) -> FineTuneResult:
     """Tokenizes text, loads the model, and runs DP-SGD fine-tuning.
 
     Args:
       rng: Random seed for batch selection and noise generation.
-      data: Sequence of ``(prompt, response)`` string pairs.
+      data: Sequence of ``(prompt, response)`` string pairs or a
+        ``pygrain.MapDataset``.
 
     Returns:
       A ``FineTuneResult`` with the trained model and merged LoRA parameters.
       Private training state (noise, optimizer) is not exposed.
     """
-    dataset = model.tokenize_texts(
-        data,
-        model_variant=self.model_variant,
-        max_seq_length=self.max_seq_length,
-    )
+    if isinstance(data, pygrain.MapDataset):
+      tokenizer = self.model_variant.tokenizer_class()
+      tokenized = data.map(
+          lambda ex: model.tokenize_example(ex, tokenizer, self.max_seq_length)
+      )
+      # before DPTrainer.
+      elements = list(tokenized)
+      dataset = {
+          'input_tokens': np.stack([x['input_tokens'] for x in elements]),
+          'loss_mask': np.stack([x['loss_mask'] for x in elements]),
+      }
+    else:
+      dataset = model.tokenize_texts(
+          data,
+          model_variant=self.model_variant,
+          max_seq_length=self.max_seq_length,
+      )
 
     lora_config = model.LoraConfig(rank=self.lora_rank)
     module, frozen_params, trainable_params = model.load_gemma(

dpsynth/text/model.py

@@ -197,6 +197,40 @@ def format_response(response: str, tokenizer: Any) -> str:
   return f'{response}{eot}'
 
 
+def tokenize_example(
+    example: tuple[str, str],
+    tokenizer: Any,
+    max_seq_length: int,
+) -> dict[str, np.ndarray]:
+  """Tokenizes a single (prompt, response) pair for supervised fine-tuning.
+
+  Prompt tokens are masked out (``loss_mask=0``) so only the response
+  contributes to the training loss.
+
+  Args:
+    example: ``(prompt, response)`` string pair.
+    tokenizer: Gemma tokenizer instance.
+    max_seq_length: Maximum sequence length (including special tokens).
+
+  Returns:
+    Dict with ``'input_tokens'`` and ``'loss_mask'`` (int32 ``[L]``).
+  """
+  prompt, response = example
+  prompt_str = format_prompt(prompt, tokenizer)
+  response_str = format_response(response, tokenizer)
+  prompt_ids = tokenizer.encode(prompt_str, add_bos=True)
+  response_ids = tokenizer.encode(response_str, add_eos=True)
+
+  ids = prompt_ids + response_ids
+  length = min(len(ids), max_seq_length)
+  tokens = np.zeros(max_seq_length, dtype=np.int32)
+  mask = np.zeros(max_seq_length, dtype=np.int32)
+  tokens[:length] = ids[:length]
+  mask[min(len(prompt_ids), length) : length] = 1
+
+  return {'input_tokens': tokens, 'loss_mask': mask}
+
+
 def tokenize_texts(
     examples: Sequence[tuple[str, str]],
     model_variant: GemmaModel,
@@ -221,20 +255,10 @@ def tokenize_texts(
   tokens = np.zeros((len(examples), max_seq_length), dtype=np.int32)
   mask = np.zeros((len(examples), max_seq_length), dtype=np.int32)
 
-  for i, (prompt, response) in enumerate(examples):
-    # Embed turn tags as strings so SentencePiece handles tokenization
-    # boundaries correctly (encoding pieces separately can shift BPE merges).
-    prompt_str = format_prompt(prompt, tokenizer)
-    response_str = format_response(response, tokenizer)
-    prompt_ids = tokenizer.encode(prompt_str, add_bos=True)
-    response_ids = tokenizer.encode(response_str, add_eos=True)
-
-    ids = prompt_ids + response_ids
-    length = min(len(ids), max_seq_length)
-    tokens[i, :length] = ids[:length]
-    # Mask: 0 for prompt, 1 for response.
-    resp_start = min(len(prompt_ids), length)
-    mask[i, resp_start:length] = 1
+  for i, example in enumerate(examples):
+    tokenized = tokenize_example(example, tokenizer, max_seq_length)
+    tokens[i] = tokenized['input_tokens']
+    mask[i] = tokenized['loss_mask']
 
   logging.info(
       'Tokenized %d examples (max_seq_length=%d)',

tests/text/dp_sft_test.py

@@ -20,13 +20,18 @@
 
 import dataclasses
 import math
+from unittest import mock
 
 from absl.testing import absltest
 from dpsynth.text import dp_sft
+from dpsynth.text import dp_trainer
 from dpsynth.text import model
+from gemma import peft
+import grain.python as pygrain
 import jax
 import jax.numpy as jnp
 from jax_privacy import execution_plan
+import numpy as np
 
 
 def _default_config():
@@ -193,6 +198,172 @@ def test_dp_event_after_calibration(self):
     event = mechanism.dp_event
     self.assertIsNotNone(event)
 
+  @mock.patch.object(peft, 'merge_params', autospec=True)
+  @mock.patch.object(dp_trainer, 'DPTrainer', autospec=True)
+  @mock.patch.object(model, 'load_gemma', autospec=True)
+  def test_call_with_map_dataset_materializes_to_numpy(
+      self, mock_load, mock_trainer_cls, mock_merge
+  ):
+    mock_load.return_value = (mock.MagicMock(), {}, {})
+    mock_trainer = mock_trainer_cls.return_value
+    mock_trainer.return_value = mock.MagicMock(params={})
+    mock_merge.return_value = {}
+
+    variant = model.GemmaModel(
+        model_class=mock.MagicMock(),
+        checkpoint_path='/mock/path',
+        tokenizer_class=_MockTokenizer,
+    )
+    pairs = [('prompt1', 'resp1'), ('prompt2', 'resp2')]
+    ds = pygrain.MapDataset.source(pairs)
+
+    fine_tuner = dp_sft.DPFineTuner(
+        model_variant=variant,
+        mechanism_config=_default_config(),
+        max_seq_length=16,
+    ).configure(zcdp_rho=0.5)
+
+    res = fine_tuner(rng=0, data=ds)
+    self.assertIsInstance(res, dp_sft.FineTuneResult)
+    mock_trainer.assert_called_once()
+    passed_dataset = mock_trainer.call_args.kwargs['data']
+    self.assertIsInstance(passed_dataset['input_tokens'], np.ndarray)
+    self.assertIsInstance(passed_dataset['loss_mask'], np.ndarray)
+    self.assertEqual(passed_dataset['input_tokens'].shape, (2, 16))
+    self.assertEqual(passed_dataset['loss_mask'].shape, (2, 16))
+
+    expected_dataset = model.tokenize_texts(pairs, variant, max_seq_length=16)
+    np.testing.assert_array_equal(
+        passed_dataset['input_tokens'], expected_dataset['input_tokens']
+    )
+    np.testing.assert_array_equal(
+        passed_dataset['loss_mask'], expected_dataset['loss_mask']
+    )
+
+  @mock.patch.object(peft, 'merge_params', autospec=True)
+  @mock.patch.object(dp_trainer, 'DPTrainer', autospec=True)
+  @mock.patch.object(model, 'load_gemma', autospec=True)
+  def test_call_with_sequence(self, mock_load, mock_trainer_cls, mock_merge):
+    mock_load.return_value = (mock.MagicMock(), {}, {})
+    mock_trainer = mock_trainer_cls.return_value
+    mock_trainer.return_value = mock.MagicMock(params={})
+    mock_merge.return_value = {}
+
+    variant = model.GemmaModel(
+        model_class=mock.MagicMock(),
+        checkpoint_path='/mock/path',
+        tokenizer_class=_MockTokenizer,
+    )
+    pairs = [('prompt1', 'resp1'), ('prompt2', 'resp2')]
+    fine_tuner = dp_sft.DPFineTuner(
+        model_variant=variant,
+        mechanism_config=_default_config(),
+        max_seq_length=16,
+    ).configure(zcdp_rho=0.5)
+
+    res = fine_tuner(rng=0, data=pairs)
+    self.assertIsInstance(res, dp_sft.FineTuneResult)
+    mock_trainer.assert_called_once()
+    passed_dataset = mock_trainer.call_args.kwargs['data']
+    expected_dataset = model.tokenize_texts(pairs, variant, max_seq_length=16)
+    np.testing.assert_array_equal(
+        passed_dataset['input_tokens'], expected_dataset['input_tokens']
+    )
+    np.testing.assert_array_equal(
+        passed_dataset['loss_mask'], expected_dataset['loss_mask']
+    )
+
+
+class _MockSpecialTokens:
+  START_OF_TURN = 0
+  END_OF_TURN = 1
+
+
+class _MockTokenizer:
+
+  def __init__(self):
+    self.special_tokens = _MockSpecialTokens()
+    self.tokens = ['<start>', '<end>']
+
+  def encode(self, text, add_bos=False, add_eos=False):
+    del add_bos, add_eos
+    return [len(text) % 10 + 1, len(text) % 5 + 1]
+
+
+class TokenizeExampleTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.tokenizer = _MockTokenizer()
+
+  def test_tokenize_example_shapes_and_types(self):
+    example = ('Hello', 'World')
+    result = model.tokenize_example(example, self.tokenizer, max_seq_length=16)
+    self.assertIn('input_tokens', result)
+    self.assertIn('loss_mask', result)
+    self.assertEqual(result['input_tokens'].shape, (16,))
+    self.assertEqual(result['loss_mask'].shape, (16,))
+    self.assertEqual(result['input_tokens'].dtype, np.int32)
+    self.assertEqual(result['loss_mask'].dtype, np.int32)
+
+  def test_tokenize_example_masking(self):
+    example = ('Hi', 'There')
+    result = model.tokenize_example(example, self.tokenizer, max_seq_length=8)
+    np.testing.assert_array_equal(result['loss_mask'][:2], [0, 0])
+    np.testing.assert_array_equal(result['loss_mask'][2:4], [1, 1])
+    np.testing.assert_array_equal(result['loss_mask'][4:], [0, 0, 0, 0])
+
+  def test_tokenize_example_truncation(self):
+    example = ('Hi', 'There')
+    result = model.tokenize_example(example, self.tokenizer, max_seq_length=3)
+    self.assertEqual(result['input_tokens'].shape, (3,))
+    self.assertEqual(result['loss_mask'].shape, (3,))
+    np.testing.assert_array_equal(result['loss_mask'], [0, 0, 1])
+
+
+class TokenizeTextsTest(absltest.TestCase):
+
+  def test_parity_with_tokenize_example(self):
+    variant = model.GemmaModel(
+        model_class=mock.MagicMock(),
+        checkpoint_path='/mock/path',
+        tokenizer_class=_MockTokenizer,
+    )
+    tokenizer = variant.tokenizer_class()
+    pairs = [('Hello', 'World'), ('How are you?', 'I am fine.')]
+    seq_result = model.tokenize_texts(
+        pairs, model_variant=variant, max_seq_length=16
+    )
+    for i, pair in enumerate(pairs):
+      ex_result = model.tokenize_example(pair, tokenizer, max_seq_length=16)
+      np.testing.assert_array_equal(
+          seq_result['input_tokens'][i], ex_result['input_tokens']
+      )
+      np.testing.assert_array_equal(
+          seq_result['loss_mask'][i], ex_result['loss_mask']
+      )
+
+  def test_parity_between_sequence_and_numpy_array(self):
+    variant = model.GemmaModel(
+        model_class=mock.MagicMock(),
+        checkpoint_path='/mock/path',
+        tokenizer_class=_MockTokenizer,
+    )
+    pairs = [('Hello', 'World'), ('How are you?', 'I am fine.')]
+    seq_result = model.tokenize_texts(
+        pairs, model_variant=variant, max_seq_length=16
+    )
+    np_result = model.tokenize_texts(
+        np.asarray(pairs), model_variant=variant, max_seq_length=16
+    )
+
+    np.testing.assert_array_equal(
+        seq_result['input_tokens'], np_result['input_tokens']
+    )
+    np.testing.assert_array_equal(
+        seq_result['loss_mask'], np_result['loss_mask']
+    )
+
 
 if __name__ == '__main__':
   absltest.main()