Method map

gemma/gm/data/_tasks.py:212–269

(self, element)

  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 to match the expected dialog template.
    # TODO(epot): Move this in a separate FormatDialog transform.
    prompt = _template.PROMPT.format(prompt)
    chosen = _template.ANSWER.format(chosen)
    rejected = _template.ANSWER.format(rejected)

    # 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