Spaces:
Paused
Paused
File size: 4,524 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | # 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.
from absl.testing import absltest
from gemma.diffusion import _chat_sampler
from gemma.diffusion import _early_stopping
from gemma.diffusion import _models
from gemma.diffusion import _sampler
from gemma.diffusion import _transformer
from gemma.gm.nn.gemma4 import _config
from gemma.gm.nn.gemma4 import _modules
_SMALL_CONFIG = _config.TransformerConfig(
num_embed=32,
embed_dim=8,
num_heads=2,
num_kv_heads=1,
head_dim=4,
hidden_dim=16,
attention_types=[_modules.AttentionType.GLOBAL],
kv_cache_sharing_config=None,
use_post_attn_norm=True,
use_post_ffw_norm=True,
final_logit_softcap=None,
global_rope_proportion=1.0,
)
_SMALL_SC_CONFIG = _transformer.SelfConditioningConfig(
features=_SMALL_CONFIG.embed_dim,
hidden_dim=_SMALL_CONFIG.hidden_dim,
)
class InterfacesTest(absltest.TestCase):
def test_chat_sampler_instantiation(self):
"""Tests that ChatSampler can be instantiated with minimal args.
Relying on default values for instantiation ensures that the user is not
overwhelmed with too many required arguments.
"""
model = _models.DiffusionGemma_26B_A4B(
config=_SMALL_CONFIG,
self_conditioning_config=_SMALL_SC_CONFIG,
)
chat = _chat_sampler.ChatSampler(model=model, params={})
# Diffusion-specific fields
self.assertEqual(chat.sampler.canvas_length, 256)
self.assertEqual(chat.sampler.max_denoising_steps, 48)
self.assertIsInstance(
chat.sampler.diffusion_process, _sampler.DiffusionProcess
)
self.assertIsInstance(
chat.sampler.diffusion_process.noise_schedule, _sampler.LinearSchedule
)
self.assertIsInstance(
chat.sampler.logit_shaper, _sampler.AnnealingTemperatureShaper
)
self.assertEqual(chat.sampler.logit_shaper.config.exponent, 1.0)
self.assertEqual(chat.sampler.logit_shaper.config.max_temperature, 0.8)
self.assertEqual(chat.sampler.logit_shaper.config.min_temperature, 0.4)
# Diffusion sampling: confidence-based selection with entropy_bound=0.1.
self.assertIsInstance(
chat.sampler.sample_from_predictions,
_sampler.SampleFromPredictions,
)
self.assertEqual(chat.sampler.sample_from_predictions.entropy_bound, 0.1)
self.assertIsInstance(
chat.sampler.early_stop_fn, _early_stopping.ChainedEarlyStop
)
self.assertLen(chat.sampler.early_stop_fn.early_stop_fns, 2)
self.assertIsInstance(
chat.sampler.early_stop_fn.early_stop_fns[0],
_early_stopping.TokenStabilityEarlyStop,
)
self.assertIsInstance(
chat.sampler.early_stop_fn.early_stop_fns[1],
_early_stopping.EntropyEarlyStop,
)
self.assertEqual(
chat.sampler.early_stop_fn.early_stop_fns[1].entropy_threshold,
0.005,
)
def test_chat_sampler_instantiation_overrides(self):
"""Tests that ChatSampler can be instantiated with optional args."""
model = _models.DiffusionGemma_26B_A4B(
config=_SMALL_CONFIG,
self_conditioning_config=_SMALL_SC_CONFIG,
)
chat = _chat_sampler.ChatSampler(
model=model,
params={},
canvas_length=32,
max_denoising_steps=16,
logit_shaper=_sampler.AnnealingTemperatureShaper(
config=_sampler.AnnealingTemperatureShaperConfig(
max_temperature=0.9,
min_temperature=0.3,
exponent=2.0,
)
),
)
# Confirm that the overrides are applied (i.e. they are different from the
# defaults)
self.assertEqual(chat.sampler.canvas_length, 32)
self.assertEqual(chat.sampler.max_denoising_steps, 16)
self.assertEqual(chat.sampler.logit_shaper.config.max_temperature, 0.9)
self.assertEqual(chat.sampler.logit_shaper.config.min_temperature, 0.3)
self.assertEqual(chat.sampler.logit_shaper.config.exponent, 2.0)
if __name__ == "__main__":
absltest.main()
|