Spaces:
Paused
Paused
File size: 3,576 Bytes
2f86d9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | # 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.
"""Base classes for host-side detokenizing metrics."""
from __future__ import annotations
import abc
import dataclasses
import functools
import flax.struct
import jax.numpy as jnp
from kauldron import kd
import numpy as np
import seqio
_GEMMA4_TOKENIZER_PATH = (
"gs://gemma-data/tokenizers/tokenizer_gemma4.model"
)
@dataclasses.dataclass(kw_only=True, frozen=True)
class BaseTokenizerMetric(kd.metrics.Metric, abc.ABC):
"""Base class for metrics requiring host-side tokenizer access."""
tokenizer_path: str = _GEMMA4_TOKENIZER_PATH
@functools.cached_property
def tokenizer(self) -> seqio.vocabularies.SentencePieceVocabulary:
return seqio.vocabularies.SentencePieceVocabulary(self.tokenizer_path)
def _squeeze_3d(self, x: jnp.ndarray) -> jnp.ndarray:
"""Squeezes 3D [B, L, 1] arrays to 2D [B, L]."""
if x.ndim == 3:
return x[..., 0]
return x
def decode_batch(self, tokens_np: np.ndarray) -> list[str]:
"""Decodes 2D token array to a list of strings on host."""
return [self.tokenizer.decode(row.tolist()) for row in tokens_np]
@dataclasses.dataclass(kw_only=True, frozen=True)
class BaseTextMetric(BaseTokenizerMetric):
"""Base class for metrics requiring host-side detokenization."""
tokens: kd.kontext.Key = kd.kontext.REQUIRED
ground_truth: kd.kontext.Key = kd.kontext.REQUIRED
@dataclasses.dataclass(kw_only=True, frozen=True)
class BaseSimpleTextMetric(BaseTextMetric):
"""Metric for simple text tasks returning a single value per example.
Stores raw token arrays in the metric state and performs all
detokenization and scoring on the host in ``compute()``.
"""
@flax.struct.dataclass
class State(kd.metrics.AutoState["BaseSimpleTextMetric"]):
"""Collects raw tokens across batches for host-side scoring."""
tokens: jnp.ndarray = kd.metrics.concat_field()
ground_truth: jnp.ndarray = kd.metrics.concat_field()
def compute(self) -> float:
"""Detokenize and score all accumulated examples on host."""
final = self.finalize()
tokens_np = np.asarray(final.tokens)
gt_np = np.asarray(final.ground_truth)
gen_texts = self.parent.decode_batch(tokens_np)
gt_texts = self.parent.decode_batch(gt_np)
num_examples = len(gen_texts)
if num_examples == 0:
return 0.0
total = sum(
float(self.parent.score_example(gen_texts[i], gt_texts[i]))
for i in range(num_examples)
)
return total / num_examples
@abc.abstractmethod
def score_example(
self, generated_text: str, ground_truth_text: str
) -> float | bool:
"""Scores a single example."""
pass
def get_state(
self,
*,
tokens: jnp.ndarray,
ground_truth: jnp.ndarray,
) -> BaseSimpleTextMetric.State:
"""Store raw tokens — all scoring is deferred to ``compute()``."""
return self.State(
tokens=self._squeeze_3d(tokens),
ground_truth=self._squeeze_3d(ground_truth),
)
|