Spaces:

Jit89
/

bmap-gemma-engine

Paused

bmap-gemma-engine / gemma /gm /data /_tasks.py

feat: initial commit for B Map AI backend and Gemma LLM service

2f86d9b 10 days ago

10.1 kB

# Copyright 2026 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""End-to-end transformations."""

from __future__ import annotations

import dataclasses

import einops
from etils.etree import jax as etree  # pylint: disable=g-importing-member
from gemma.gm.data import _functional
from gemma.gm.text import _tokenizer
from grain import python as grain
import jax
from kauldron import kd
import numpy as np

# Turn tag strings indexed by tokenizer FORMAT.
# Gemma4 uses '<|turn>' / '<turn|>', all others use '<start_of_turn>' /
# '<end_of_turn>'. Importing the `dialog` library is intentionally avoided
# to keep the dep footprint small; the two known format strings are inlined.
_TURN_TAGS: dict[str, tuple[str, str]] = {}


def _get_turn_tags(
    tokenizer: _tokenizer.Tokenizer,
) -> tuple[str, str]:
  """Returns (start_of_turn, end_of_turn) tag strings for *tokenizer*."""
  fmt = getattr(tokenizer, 'FORMAT', None)
  # dialog.Format.GEMMA4 has value 'gemma4' (StrEnum).
  if fmt is not None and str(fmt).lower() == 'gemma4':
    return ('<|turn>', '<turn|>')
  # Default: Gemma3 / legacy format.
  return ('<start_of_turn>', '<end_of_turn>')


def _format_prompt(prompt: str, tokenizer: _tokenizer.Tokenizer) -> str:
  """Formats *prompt* with the correct turn tags for *tokenizer*."""
  sot, eot = _get_turn_tags(tokenizer)
  return f'{sot}user\n{prompt}{eot}\n{sot}model\n'


def _format_answer(response: str, tokenizer: _tokenizer.Tokenizer) -> str:
  """Formats *response* with the correct turn tags for *tokenizer*."""
  _, eot = _get_turn_tags(tokenizer)
  return f'{response}{eot}'


@dataclasses.dataclass(kw_only=True, frozen=True)
class Seq2SeqTask(grain.MapTransform):
  """Sequence-to-sequence task.
  This task will:
  * Format the prompt and response to match the expected dialog template (i.e.
    add the `<start_of_turn>user`, `<end_of_turn>`,...)
  * Tokenize the prompt and response.
  * Concatenate the input and response to create the model input and target
    (target is the input shifted by one token).
  * Create the loss mask (0 for prompt, 1 for response)
  * Pad/truncate the input and target to the max length.
  Example:
  ```python
  # Input:
  {
      'prompt': 'Hello! I would love to visit France.',
      'response': 'Bonjour ! J'adorerais visiter la France.',
  }
  # Ouptut:
  {
      'input': i32['max_length'],
      'target': i32['max_length 1'],
      'target_mask': bool['max_length 1'],
  }
  ```
  Note:
    - Input and target are the same sequence shifted by one token.
    - The last token from the target is truncated from the input (as there's no
      target for it)
  Attributes:
    in_prompt: Input key
    in_response: Input key
    out_input: Output key (will be added to the example dict)
    out_target: Output key (will be added to the example dict)
    out_target_mask: Output key (will be added to the example dict)
    drop_inputs: If True, remove the input keys from the output.
    max_length: The max length of the sequence (examples will be
      padded/truncated to this length).
    truncate: Whether to truncate the sequence to the max length. If `False`,
      sequences longer than the `max_length` will raise an error.
    sampling: If `True`, the dataset will yield the original prompt and response
      so they can be used inside `gm.evals.SamplerEvaluator`.
  """

  in_prompt: kd.kontext.Key  # e.g. `'prompt'`
  in_response: kd.kontext.Key  # e.g. `'response'`

  out_input: kd.kontext.Key  # e.g. `'input'`
  out_target: kd.kontext.Key  # e.g. `'target'`
  out_target_mask: kd.kontext.Key  # e.g. `'target_mask'`

  drop_inputs: bool = True

  tokenizer: _tokenizer.Tokenizer

  # Padding parameters
  max_length: int
  truncate: bool = False

  sampling: bool = False

  def map(self, element):
    # Deep-copy to avoid mutating the input.
    element = etree.copy(element)

    # Extract the values from the `dict` example.
    # `kontext.get_by_path(element, self.in_prompt)` is equivalent to
    # `element[self.in_prompt]`, but supports nested dicts and dataclasses.
    prompt = kd.kontext.get_by_path(element, self.in_prompt)
    response = kd.kontext.get_by_path(element, self.in_response)

    # TODO(epot): Supports nested drop
    if self.drop_inputs:
      del element[self.in_prompt]  # pyrefly: ignore[unsupported-operation]
      del element[self.in_response]  # pyrefly: ignore[unsupported-operation]

    # Some datasets (TFDS) returns `bytes` instead of `str`, so decode them.
    prompt = _decode_bytes(prompt)
    response = _decode_bytes(response)

    # Format the input using tokenizer-aware turn tags.
    # TODO(epot): Add a `template` protocol for full customization.
    prompt = _format_prompt(prompt, self.tokenizer)
    response = _format_answer(response, self.tokenizer)

    # For sampling, we don't need to tokenize the input.
    if self.sampling:
      kd.kontext.set_by_path(element, self.out_input, prompt)
      kd.kontext.set_by_path(element, self.out_target, response)
      return element

    # Tokenize the input and the responses.
    prompt = self.tokenizer.encode(prompt, add_bos=True)
    response = self.tokenizer.encode(response)

    # Create the model inputs/targets/loss_mask.
    out = _functional.make_seq2seq_fields(
        prompt=prompt,
        response=response,
    )

    # Add padding.
    out = _functional.pad(
        out,
        max_length=self.max_length,
        truncate=self.truncate,
    )

    # For shape compatibility with the loss
    target = einops.rearrange(out.target, '... -> ... 1')
    target_mask = einops.rearrange(out.target_mask, '... -> ... 1')

    # Add the fields to the output `dict`.
    # Equivalent to `element[self.out_input] = ...`
    kd.kontext.set_by_path(element, self.out_input, out.input)
    kd.kontext.set_by_path(element, self.out_target, target)
    kd.kontext.set_by_path(element, self.out_target_mask, target_mask)
    return element


@dataclasses.dataclass(kw_only=True, frozen=True)
class ContrastiveTask(grain.MapTransform):
  """Creates the contrastive model inputs for DPO-like loss.
  Input:
  ```python
  {
      'prompt': 'How much are 2+2 ?',
      'chosen': 'Yes, this is 4.',
      'rejected': 'Of course, 2+2 is 42.',
  }
  ```
  Output:
  ```python
  {
      'tokens': i32['2 max_length'],
      'mask': bool['2 max_length'],
  }
  ```
  In the output, `[chosen, rejected]` token ids are stacked (in that order).
  Attributes:
    in_prompt: Input key
    in_chosen: Input key
    in_rejected: Input key
    out_tokens: Output key (will be added to the example dict)
    out_mask: Output key (will be added to the example dict)
    tokenizer: The tokenizer to use.
    max_length: The max length of the sequence (examples will be
      padded/truncated to this length).
    truncate: Whether to truncate the sequence to the max length. If `False`,
      sequences longer than the `max_length` will raise an error.
    drop_inputs: If True, remove the input keys from the output.
  """

  in_prompt: kd.kontext.Key  # e.g. `'input'`
  in_chosen: kd.kontext.Key  # e.g. `'chosen'`
  in_rejected: kd.kontext.Key  # e.g. `'rejected'`

  out_tokens: kd.kontext.Key  # e.g. `'tokens'`
  out_targets: kd.kontext.Key  # e.g. `'target'`
  out_mask: kd.kontext.Key  # e.g. `'mask'`

  tokenizer: _tokenizer.Tokenizer

  max_length: int
  truncate: bool = False

  drop_inputs: bool = True

  def map(self, element):
    prompt = kd.kontext.get_by_path(element, self.in_prompt)
    chosen = kd.kontext.get_by_path(element, self.in_chosen)
    rejected = kd.kontext.get_by_path(element, self.in_rejected)

    # Some datasets (TFDS) returns `bytes` instead of `str`, so decode them.
    prompt = _decode_bytes(prompt)
    chosen = _decode_bytes(chosen)
    rejected = _decode_bytes(rejected)

    # Format the input using tokenizer-aware turn tags.
    # TODO(epot): Extract into a standalone FormatDialog transform.
    prompt = _format_prompt(prompt, self.tokenizer)
    chosen = _format_answer(chosen, self.tokenizer)
    rejected = _format_answer(rejected, self.tokenizer)

    # Tokenize the input and the responses.
    # Note: Input should start with begin-of-sequence token.
    prompt = self.tokenizer.encode(prompt, add_bos=True)
    chosen = self.tokenizer.encode(chosen)
    rejected = self.tokenizer.encode(rejected)

    next_token_chosen = _functional.make_seq2seq_fields(
        prompt=prompt,
        response=chosen,
    )
    next_token_rejected = _functional.make_seq2seq_fields(
        prompt=prompt,
        response=rejected,
    )

    # Add padding.
    (next_token_chosen, next_token_rejected) = _functional.pad(
        (next_token_chosen, next_token_rejected),
        max_length=self.max_length,
        truncate=self.truncate,
    )

    # Stack the input and target.
    out = jax.tree.map(
        lambda x, y: np.stack([x, y], axis=0),
        next_token_chosen,
        next_token_rejected,
    )

    # Add the fields to the output `dict`.
    # Equivalent to `element[self.out_input] = ...`
    kd.kontext.set_by_path(element, self.out_tokens, out.input)
    kd.kontext.set_by_path(element, self.out_targets, out.target)
    kd.kontext.set_by_path(element, self.out_mask, out.target_mask)

    # TODO(epot): Supports nested drop
    if self.drop_inputs:
      del element[self.in_prompt]
      del element[self.in_chosen]
      del element[self.in_rejected]

    return element


def _decode_bytes(element):
  if isinstance(element, bytes):
    return element.decode('utf-8')
  else:
    return element

	# Copyright 2026 DeepMind Technologies Limited.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""End-to-end transformations."""

	from __future__ import annotations

	import dataclasses

	import einops
	from etils.etree import jax as etree # pylint: disable=g-importing-member
	from gemma.gm.data import _functional
	from gemma.gm.text import _tokenizer
	from grain import python as grain
	import jax
	from kauldron import kd
	import numpy as np

	# Turn tag strings indexed by tokenizer FORMAT.
	# Gemma4 uses '<\|turn>' / '<turn\|>', all others use '<start_of_turn>' /
	# '<end_of_turn>'. Importing the `dialog` library is intentionally avoided
	# to keep the dep footprint small; the two known format strings are inlined.
	_TURN_TAGS: dict[str, tuple[str, str]] = {}


	def _get_turn_tags(
	tokenizer: _tokenizer.Tokenizer,
	) -> tuple[str, str]:
	"""Returns (start_of_turn, end_of_turn) tag strings for tokenizer."""
	fmt = getattr(tokenizer, 'FORMAT', None)
	# dialog.Format.GEMMA4 has value 'gemma4' (StrEnum).
	if fmt is not None and str(fmt).lower() == 'gemma4':
	return ('<\|turn>', '<turn\|>')
	# Default: Gemma3 / legacy format.
	return ('<start_of_turn>', '<end_of_turn>')


	def _format_prompt(prompt: str, tokenizer: _tokenizer.Tokenizer) -> str:
	"""Formats prompt with the correct turn tags for tokenizer."""
	sot, eot = _get_turn_tags(tokenizer)
	return f'{sot}user\n{prompt}{eot}\n{sot}model\n'


	def _format_answer(response: str, tokenizer: _tokenizer.Tokenizer) -> str:
	"""Formats response with the correct turn tags for tokenizer."""
	_, eot = _get_turn_tags(tokenizer)
	return f'{response}{eot}'


	@dataclasses.dataclass(kw_only=True, frozen=True)
	class Seq2SeqTask(grain.MapTransform):
	"""Sequence-to-sequence task.

	This task will:

	* Format the prompt and response to match the expected dialog template (i.e.
	add the `<start_of_turn>user`, `<end_of_turn>`,...)
	* Tokenize the prompt and response.
	* Concatenate the input and response to create the model input and target
	(target is the input shifted by one token).
	* Create the loss mask (0 for prompt, 1 for response)
	* Pad/truncate the input and target to the max length.

	Example:

	```python
	# Input:
	{
	'prompt': 'Hello! I would love to visit France.',
	'response': 'Bonjour ! J'adorerais visiter la France.',
	}
	# Ouptut:
	{
	'input': i32['max_length'],
	'target': i32['max_length 1'],
	'target_mask': bool['max_length 1'],
	}
	```

	Note:
	- Input and target are the same sequence shifted by one token.
	- The last token from the target is truncated from the input (as there's no
	target for it)

	Attributes:
	in_prompt: Input key
	in_response: Input key
	out_input: Output key (will be added to the example dict)
	out_target: Output key (will be added to the example dict)
	out_target_mask: Output key (will be added to the example dict)
	drop_inputs: If True, remove the input keys from the output.
	max_length: The max length of the sequence (examples will be
	padded/truncated to this length).
	truncate: Whether to truncate the sequence to the max length. If `False`,
	sequences longer than the `max_length` will raise an error.
	sampling: If `True`, the dataset will yield the original prompt and response
	so they can be used inside `gm.evals.SamplerEvaluator`.
	"""

	in_prompt: kd.kontext.Key # e.g. `'prompt'`
	in_response: kd.kontext.Key # e.g. `'response'`

	out_input: kd.kontext.Key # e.g. `'input'`
	out_target: kd.kontext.Key # e.g. `'target'`
	out_target_mask: kd.kontext.Key # e.g. `'target_mask'`

	drop_inputs: bool = True

	tokenizer: _tokenizer.Tokenizer

	# Padding parameters
	max_length: int
	truncate: bool = False

	sampling: bool = False

	def map(self, element):
	# Deep-copy to avoid mutating the input.
	element = etree.copy(element)

	# Extract the values from the `dict` example.
	# `kontext.get_by_path(element, self.in_prompt)` is equivalent to
	# `element[self.in_prompt]`, but supports nested dicts and dataclasses.
	prompt = kd.kontext.get_by_path(element, self.in_prompt)
	response = kd.kontext.get_by_path(element, self.in_response)

	# TODO(epot): Supports nested drop
	if self.drop_inputs:
	del element[self.in_prompt] # pyrefly: ignore[unsupported-operation]
	del element[self.in_response] # pyrefly: ignore[unsupported-operation]

	# Some datasets (TFDS) returns `bytes` instead of `str`, so decode them.
	prompt = _decode_bytes(prompt)
	response = _decode_bytes(response)

	# Format the input using tokenizer-aware turn tags.
	# TODO(epot): Add a `template` protocol for full customization.
	prompt = _format_prompt(prompt, self.tokenizer)
	response = _format_answer(response, self.tokenizer)

	# For sampling, we don't need to tokenize the input.
	if self.sampling:
	kd.kontext.set_by_path(element, self.out_input, prompt)
	kd.kontext.set_by_path(element, self.out_target, response)
	return element

	# Tokenize the input and the responses.
	prompt = self.tokenizer.encode(prompt, add_bos=True)
	response = self.tokenizer.encode(response)

	# Create the model inputs/targets/loss_mask.
	out = _functional.make_seq2seq_fields(
	prompt=prompt,
	response=response,
	)

	# Add padding.
	out = _functional.pad(
	out,
	max_length=self.max_length,
	truncate=self.truncate,
	)

	# For shape compatibility with the loss
	target = einops.rearrange(out.target, '... -> ... 1')
	target_mask = einops.rearrange(out.target_mask, '... -> ... 1')

	# Add the fields to the output `dict`.
	# Equivalent to `element[self.out_input] = ...`
	kd.kontext.set_by_path(element, self.out_input, out.input)
	kd.kontext.set_by_path(element, self.out_target, target)
	kd.kontext.set_by_path(element, self.out_target_mask, target_mask)
	return element


	@dataclasses.dataclass(kw_only=True, frozen=True)
	class ContrastiveTask(grain.MapTransform):
	"""Creates the contrastive model inputs for DPO-like loss.

	Input:

	```python
	{
	'prompt': 'How much are 2+2 ?',
	'chosen': 'Yes, this is 4.',
	'rejected': 'Of course, 2+2 is 42.',
	}
	```

	Output:

	```python
	{
	'tokens': i32['2 max_length'],
	'mask': bool['2 max_length'],
	}
	```

	In the output, `[chosen, rejected]` token ids are stacked (in that order).

	Attributes:
	in_prompt: Input key
	in_chosen: Input key
	in_rejected: Input key
	out_tokens: Output key (will be added to the example dict)
	out_mask: Output key (will be added to the example dict)
	tokenizer: The tokenizer to use.
	max_length: The max length of the sequence (examples will be
	padded/truncated to this length).
	truncate: Whether to truncate the sequence to the max length. If `False`,
	sequences longer than the `max_length` will raise an error.
	drop_inputs: If True, remove the input keys from the output.
	"""

	in_prompt: kd.kontext.Key # e.g. `'input'`
	in_chosen: kd.kontext.Key # e.g. `'chosen'`
	in_rejected: kd.kontext.Key # e.g. `'rejected'`

	out_tokens: kd.kontext.Key # e.g. `'tokens'`
	out_targets: kd.kontext.Key # e.g. `'target'`
	out_mask: kd.kontext.Key # e.g. `'mask'`

	tokenizer: _tokenizer.Tokenizer

	max_length: int
	truncate: bool = False

	drop_inputs: bool = True

	def map(self, element):
	prompt = kd.kontext.get_by_path(element, self.in_prompt)
	chosen = kd.kontext.get_by_path(element, self.in_chosen)
	rejected = kd.kontext.get_by_path(element, self.in_rejected)

	# Some datasets (TFDS) returns `bytes` instead of `str`, so decode them.
	prompt = _decode_bytes(prompt)
	chosen = _decode_bytes(chosen)
	rejected = _decode_bytes(rejected)

	# Format the input using tokenizer-aware turn tags.
	# TODO(epot): Extract into a standalone FormatDialog transform.
	prompt = _format_prompt(prompt, self.tokenizer)
	chosen = _format_answer(chosen, self.tokenizer)
	rejected = _format_answer(rejected, self.tokenizer)

	# Tokenize the input and the responses.
	# Note: Input should start with begin-of-sequence token.
	prompt = self.tokenizer.encode(prompt, add_bos=True)
	chosen = self.tokenizer.encode(chosen)
	rejected = self.tokenizer.encode(rejected)

	next_token_chosen = _functional.make_seq2seq_fields(
	prompt=prompt,
	response=chosen,
	)
	next_token_rejected = _functional.make_seq2seq_fields(
	prompt=prompt,
	response=rejected,
	)

	# Add padding.
	(next_token_chosen, next_token_rejected) = _functional.pad(
	(next_token_chosen, next_token_rejected),
	max_length=self.max_length,
	truncate=self.truncate,
	)

	# Stack the input and target.
	out = jax.tree.map(
	lambda x, y: np.stack([x, y], axis=0),
	next_token_chosen,
	next_token_rejected,
	)

	# Add the fields to the output `dict`.
	# Equivalent to `element[self.out_input] = ...`
	kd.kontext.set_by_path(element, self.out_tokens, out.input)
	kd.kontext.set_by_path(element, self.out_targets, out.target)
	kd.kontext.set_by_path(element, self.out_mask, out.target_mask)

	# TODO(epot): Supports nested drop
	if self.drop_inputs:
	del element[self.in_prompt]
	del element[self.in_chosen]
	del element[self.in_rejected]

	return element


	def _decode_bytes(element):
	if isinstance(element, bytes):
	return element.decode('utf-8')
	else:
	return element