wehub-resource-sync
e06fe8e8c6
Self-hosted runner (benchmark) / Benchmark (aws-g5-4xlarge-cache) (push) Waiting to run
New model PR merged notification / Notify new model (push) Waiting to run
Update Transformers metadata / build_and_package (push) Waiting to run
Secret Leaks / trufflehog (push) Failing after 1s
Build documentation / build (push) Failing after 1s
Build documentation / build_other_lang (push) Failing after 0s
CodeQL Security Analysis / CodeQL Analysis (push) Failing after 0s
PR CI / pr-ci (push) Failing after 1s
Slow tests on important models (on Push - A10) / Get all modified files (push) Failing after 1s
Slow tests on important models (on Push - A10) / Model CI (push) Has been skipped
319 lines
15 KiB
Python
319 lines
15 KiB
Python
# Copyright 2026 the HuggingFace Team. All rights reserved.
|
|
#
|
|
# 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.
|
|
|
|
import tempfile
|
|
import unittest
|
|
|
|
from transformers import is_torch_available
|
|
from transformers.testing_utils import require_torch, torch_device
|
|
|
|
|
|
if is_torch_available():
|
|
import torch
|
|
|
|
from transformers import (
|
|
DiffusionGemmaGenerationConfig,
|
|
DiffusionGemmaGenerationMixin,
|
|
EntropyBoundSampler,
|
|
EntropyBoundSamplerConfig,
|
|
GenerationConfig,
|
|
LinearTemperatureScheduleLogitsProcessor,
|
|
StableAndConfidentStoppingCriteria,
|
|
)
|
|
|
|
|
|
@require_torch
|
|
class DiffusionGemmaGenerationClassesTester(unittest.TestCase):
|
|
def test_generation_config_interface(self):
|
|
"""
|
|
Test to confirm that basic `GenerationConfig` are also accepted in `DiffusionGemmaGenerationConfig`.
|
|
|
|
If this test pass, it implies that text diffusion has roughly the same interface as AR generation, since we
|
|
can pass kwargs from `generate` to the generation config.
|
|
"""
|
|
basic_parameterization = {
|
|
"max_length": 128,
|
|
"max_new_tokens": 64,
|
|
"cache_implementation": "dynamic",
|
|
"pad_token_id": 0,
|
|
"eos_token_id": 1,
|
|
}
|
|
diffusion_generation_config = DiffusionGemmaGenerationConfig(**basic_parameterization)
|
|
ar_generation_config = GenerationConfig(**basic_parameterization)
|
|
|
|
for attr in basic_parameterization.keys():
|
|
self.assertEqual(getattr(diffusion_generation_config, attr), getattr(ar_generation_config, attr))
|
|
|
|
def test_bad_diffusion_generation_config_parameterization(self):
|
|
"""
|
|
Test to ensure that we raise an error when users try to add AR-specific parameters to the
|
|
`DiffusionGemmaGenerationConfig`.
|
|
"""
|
|
# Some AR-specific parameters
|
|
ar_parameters = {
|
|
"do_sample": True,
|
|
"num_beams": 4,
|
|
"num_beam_groups": 1,
|
|
"temperature": 0.7,
|
|
"top_k": 50,
|
|
"top_p": 0.9,
|
|
"repetition_penalty": 1.0,
|
|
"no_repeat_ngram_size": 0,
|
|
"encoder_no_repeat_ngram_size": 0,
|
|
"length_penalty": 1.0,
|
|
"early_stopping": False,
|
|
"num_return_sequences": 1,
|
|
"foo": "bar", # random kwargs are also not accepted
|
|
}
|
|
# All these should raise an exception
|
|
for ar_param, value in ar_parameters.items():
|
|
with self.assertRaises(ValueError, msg=f"key={ar_param}"):
|
|
DiffusionGemmaGenerationConfig(**{ar_param: value})
|
|
|
|
def test_save_load_generation_config(self):
|
|
"""
|
|
Tests that we can save and load a DiffusionGemmaGenerationConfig, including its inner config dataclasses
|
|
(e.g. a sampler config)
|
|
"""
|
|
original_config = DiffusionGemmaGenerationConfig(
|
|
max_new_tokens=64,
|
|
sampler_config=EntropyBoundSamplerConfig(entropy_bound=0.1),
|
|
t_min=0.4,
|
|
t_max=0.8,
|
|
stability_threshold=1,
|
|
confidence_threshold=0.005,
|
|
)
|
|
test_attrs = (
|
|
"max_new_tokens",
|
|
"sampler_config",
|
|
"t_min",
|
|
"t_max",
|
|
"stability_threshold",
|
|
"confidence_threshold",
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
original_config.save_pretrained(tmp_dir)
|
|
loaded_config = DiffusionGemmaGenerationConfig.from_pretrained(tmp_dir)
|
|
for attr_name in test_attrs:
|
|
original_attr = getattr(original_config, attr_name)
|
|
loaded_attr = getattr(loaded_config, attr_name)
|
|
self.assertEqual(original_attr, loaded_attr) # same class, same contents
|
|
|
|
def test_eb_sampler_initialize_canvas(self):
|
|
"""
|
|
Tests that `initialize_canvas` is working as expected for `EntropyBoundSampler`.
|
|
Canvas ininitalization is random. Two samples are extremelly unlikely to be the same.
|
|
"""
|
|
sampler = _get_eb_sampler()
|
|
canvas_1 = sampler.initialize_canvas(batch_size=1, device=torch_device)
|
|
canvas_2 = sampler.initialize_canvas(batch_size=1, device=torch_device)
|
|
self.assertFalse((canvas_1 == canvas_2).all())
|
|
|
|
def test_eb_sampler_accept_canvas(self):
|
|
"""
|
|
Tests that `accept_canvas` is working as expected for `EntropyBoundSampler`.
|
|
Please see comments in the test for expected logic and corner cases.
|
|
"""
|
|
|
|
# very loose explanation: the `entropy-bound` (EB) variable controls how much entropy we're willing to accept
|
|
sampler_low_eb = _get_eb_sampler(entropy_bound=1e-2)
|
|
sampler_high_eb = _get_eb_sampler(entropy_bound=1e-1)
|
|
|
|
current_canvas = sampler_high_eb.initialize_canvas(batch_size=1, device=torch_device)
|
|
denoiser_canvas = sampler_high_eb.initialize_canvas(batch_size=1, device=torch_device)
|
|
|
|
# create logits such that all positions have high entropy, except for a few select cases
|
|
# NOTE: the first token above the threshold is accepted
|
|
logits = torch.zeros((1, 256, 10000), device=torch_device) # [bsz, canvas_len, vocab_size]
|
|
logits[0, 0, 0] = 1.8e1 # token entropy at position 0 = 2.9e-3 -> accepted in both cases
|
|
logits[0, 1, 1] = 1.45e1 # token entropy at position 1 = 7.8e-2 -> accepted in both cases
|
|
logits[0, 2, 2] = 1.45e1 # token entropy at position 2 = 7.8e-2 -> accepted only in the high eb case
|
|
|
|
# higher EB -> more accepted tokens
|
|
accepted_high_eb = sampler_high_eb.accept_canvas(
|
|
current_canvas=current_canvas, denoiser_canvas=denoiser_canvas, logits=logits, cur_step=None
|
|
)
|
|
accepted_low_eb = sampler_low_eb.accept_canvas(
|
|
current_canvas=current_canvas, denoiser_canvas=denoiser_canvas, logits=logits, cur_step=None
|
|
)
|
|
num_accepted_high_eb = (accepted_high_eb == denoiser_canvas).sum().item()
|
|
num_accepted_low_eb = (accepted_low_eb == denoiser_canvas).sum().item()
|
|
self.assertTrue(num_accepted_high_eb == num_accepted_low_eb + 1)
|
|
|
|
def test_eb_sampler_renoise_canvas(self):
|
|
"""
|
|
Tests that `renoise_canvas` is working as expected for `EntropyBoundSampler`.
|
|
All non-accepted tokens are renoised.
|
|
"""
|
|
sampler = _get_eb_sampler(entropy_bound=1e-1)
|
|
|
|
# NOTE: `renoise_canvas` is stateful: depends on the outcome of `accept_canvas`, as it renoises all
|
|
# non-accepted tokens
|
|
current_canvas = sampler.initialize_canvas(batch_size=1, device=torch_device)
|
|
denoiser_canvas = sampler.initialize_canvas(batch_size=1, device=torch_device)
|
|
|
|
logits = torch.zeros((1, 256, 10000), device=torch_device) # [bsz, canvas_len, vocab_size]
|
|
# corresponding token entropy = 0 -> these 9 tokens will definitely get accepted and, therefore, not renoised
|
|
# (but the first token above the threshold is also accepted, so we'll have 9+1 accepted tokens)
|
|
logits[0, :9, 0] = 1e6
|
|
|
|
accepted_canvas = sampler.accept_canvas(
|
|
current_canvas=current_canvas, denoiser_canvas=denoiser_canvas, logits=logits, cur_step=None
|
|
)
|
|
renoised_canvas = sampler.renoise_canvas(accepted_canvas=accepted_canvas, cur_step=None)
|
|
num_not_renoised_canvas = (renoised_canvas == accepted_canvas).sum().item()
|
|
self.assertGreaterEqual(num_not_renoised_canvas, 10) # can be >10 if the same token is sampled in the same pos
|
|
self.assertTrue((accepted_canvas[0, :9] == renoised_canvas[0, :9]).all())
|
|
|
|
def test_linear_temperature_schedule(self):
|
|
t_min = 0.4
|
|
t_max = 0.8
|
|
max_dns = 48
|
|
logits_processor = LinearTemperatureScheduleLogitsProcessor(
|
|
t_min=t_min, t_max=t_max, max_denoising_steps=max_dns
|
|
)
|
|
scores = torch.ones((1, 10), device=torch_device)
|
|
|
|
# cur_step == max_denoising_steps -> applies maximum temperature
|
|
modified_scores = logits_processor(input_ids=None, scores=scores, cur_step=max_dns)
|
|
self.assertTrue((modified_scores == scores / t_max).all())
|
|
|
|
# cur_step == max_denoising_steps/2 -> applies (t_max + t_min)/2
|
|
modified_scores = logits_processor(input_ids=None, scores=scores, cur_step=max_dns / 2)
|
|
self.assertTrue((modified_scores == scores / ((t_max + t_min) / 2)).all())
|
|
|
|
def test_stable_and_confident_stopping_criteria_confidence(self):
|
|
"""
|
|
Tests the behaviour of `confidence_threshold` in `StableAndConfidentStoppingCriteria`
|
|
"""
|
|
stopping_criteria_strict = StableAndConfidentStoppingCriteria(stability_threshold=0, confidence_threshold=1e-2)
|
|
# vocab size = 10000 -> max entropy = 9.21 -> a confidence threshold >9.21 will accept everything
|
|
stopping_criteria_lax = StableAndConfidentStoppingCriteria(stability_threshold=0, confidence_threshold=9.20)
|
|
stopping_criteria_too_lax = StableAndConfidentStoppingCriteria(
|
|
stability_threshold=0, confidence_threshold=9.22
|
|
)
|
|
|
|
# this should NEVER trigger the stopping criteria, assuming the the theshold is < ln(1/vocab_size)
|
|
logits_max_entropy = torch.zeros((1, 10, 10000), device=torch_device)
|
|
self.assertFalse(stopping_criteria_strict(argmax_canvas=None, logits=logits_max_entropy).all())
|
|
self.assertFalse(stopping_criteria_lax(argmax_canvas=None, logits=logits_max_entropy).all())
|
|
# # sanity-check
|
|
self.assertTrue(stopping_criteria_too_lax(argmax_canvas=None, logits=logits_max_entropy).all())
|
|
|
|
# mean entropy = 7.8e-2 -> only the lax triggers
|
|
logits_medium_entropy = torch.zeros((1, 10, 10000), device=torch_device)
|
|
logits_medium_entropy[:, :, 0] = 1.45e1
|
|
self.assertFalse(stopping_criteria_strict(argmax_canvas=None, logits=logits_medium_entropy).all())
|
|
self.assertTrue(stopping_criteria_lax(argmax_canvas=None, logits=logits_medium_entropy).all())
|
|
|
|
# mean entropy = 2.9e-3 -> both trigger
|
|
logits_low_entropy = torch.zeros((1, 10, 10000), device=torch_device)
|
|
logits_low_entropy[:, :, 0] = 1.8e1
|
|
self.assertTrue(stopping_criteria_strict(argmax_canvas=None, logits=logits_low_entropy).all())
|
|
self.assertTrue(stopping_criteria_lax(argmax_canvas=None, logits=logits_low_entropy).all())
|
|
|
|
def test_stable_and_confident_stopping_criteria_stability(self):
|
|
"""
|
|
Tests the behaviour of `stability_threshold` in `StableAndConfidentStoppingCriteria`
|
|
"""
|
|
# vocab size = 10000 -> max entropy = 9.21 -> a confidence threshold >9.21 will accept everything
|
|
stopping_criteria_1 = StableAndConfidentStoppingCriteria(stability_threshold=1, confidence_threshold=9.22)
|
|
stopping_criteria_2 = StableAndConfidentStoppingCriteria(stability_threshold=2, confidence_threshold=9.22)
|
|
|
|
logits = torch.zeros((1, 10, 10000), device=torch_device) # mean entropy = 9.21
|
|
argmax_canvas_1 = torch.randint(low=0, high=10000, size=(1, 10), device=torch_device)
|
|
argmax_canvas_2 = torch.randint(low=0, high=10000, size=(1, 10), device=torch_device)
|
|
|
|
# In both cases, they won't trigger after 1 canvas (needs to meet the stability criteria)
|
|
self.assertFalse(stopping_criteria_1(argmax_canvas=argmax_canvas_1, logits=logits).all())
|
|
self.assertFalse(stopping_criteria_2(argmax_canvas=argmax_canvas_1, logits=logits).all())
|
|
|
|
# `stopping_criteria_1` will be happy after 2 steps with the same canvas
|
|
self.assertTrue(stopping_criteria_1(argmax_canvas=argmax_canvas_1, logits=logits).all())
|
|
self.assertFalse(stopping_criteria_2(argmax_canvas=argmax_canvas_1, logits=logits).all())
|
|
|
|
# both will be happy after 3 steps with the same canvas
|
|
self.assertTrue(stopping_criteria_1(argmax_canvas=argmax_canvas_1, logits=logits).all())
|
|
self.assertTrue(stopping_criteria_2(argmax_canvas=argmax_canvas_1, logits=logits).all())
|
|
|
|
# If we pass a different canvas, the stability criteria will be set to false
|
|
self.assertFalse(stopping_criteria_1(argmax_canvas=argmax_canvas_2, logits=logits).all())
|
|
self.assertFalse(stopping_criteria_2(argmax_canvas=argmax_canvas_2, logits=logits).all())
|
|
|
|
def test_tokens_per_forward(self):
|
|
"""
|
|
Tests that the tokens per forward implementation is working as expected, for bsz == 1
|
|
"""
|
|
input_ids = torch.tensor([[5] * 100], dtype=torch.int)
|
|
decoder_forward_passes = torch.tensor([10], dtype=torch.int)
|
|
initial_input_ids_len = 0
|
|
pad_token_id = 1
|
|
|
|
tokens_per_forward = DiffusionGemmaGenerationMixin._compute_tokens_per_forward(
|
|
input_ids, decoder_forward_passes, initial_input_ids_len, pad_token_id
|
|
)
|
|
self.assertEqual(tokens_per_forward[0], 100 / 10)
|
|
|
|
initial_input_ids_len = 10
|
|
tokens_per_forward = DiffusionGemmaGenerationMixin._compute_tokens_per_forward(
|
|
input_ids, decoder_forward_passes, initial_input_ids_len, pad_token_id
|
|
)
|
|
self.assertEqual(tokens_per_forward[0], (100 - 10) / 10)
|
|
|
|
input_ids[:, -30:] = pad_token_id
|
|
tokens_per_forward = DiffusionGemmaGenerationMixin._compute_tokens_per_forward(
|
|
input_ids, decoder_forward_passes, initial_input_ids_len, pad_token_id
|
|
)
|
|
self.assertEqual(tokens_per_forward[0], (100 - 10 - 30) / 10)
|
|
|
|
def test_tokens_per_forward_batched(self):
|
|
"""
|
|
Tests that the tokens per forward implementation is working as expected, for bsz > 1
|
|
"""
|
|
input_ids = torch.tensor([[5] * 100] * 2, dtype=torch.int)
|
|
decoder_forward_passes = torch.tensor([10, 7], dtype=torch.int)
|
|
initial_input_ids_len = 0
|
|
pad_token_id = 1
|
|
|
|
tokens_per_forward = DiffusionGemmaGenerationMixin._compute_tokens_per_forward(
|
|
input_ids, decoder_forward_passes, initial_input_ids_len, pad_token_id
|
|
)
|
|
torch.testing.assert_close(tokens_per_forward, torch.tensor([100 / 10, 100 / 7]))
|
|
|
|
initial_input_ids_len = 10
|
|
tokens_per_forward = DiffusionGemmaGenerationMixin._compute_tokens_per_forward(
|
|
input_ids, decoder_forward_passes, initial_input_ids_len, pad_token_id
|
|
)
|
|
torch.testing.assert_close(tokens_per_forward, torch.tensor([(100 - 10) / 10, (100 - 10) / 7]))
|
|
|
|
input_ids[0, -30:] = pad_token_id
|
|
input_ids[1, -15:] = pad_token_id
|
|
tokens_per_forward = DiffusionGemmaGenerationMixin._compute_tokens_per_forward(
|
|
input_ids, decoder_forward_passes, initial_input_ids_len, pad_token_id
|
|
)
|
|
torch.testing.assert_close(tokens_per_forward, torch.tensor([(100 - 10 - 30) / 10, (100 - 10 - 15) / 7]))
|
|
|
|
|
|
def _get_eb_sampler(entropy_bound: float = 0.1) -> EntropyBoundSampler:
|
|
"""Returns a parameterized `EntropyBoundSampler`"""
|
|
sampler_config = EntropyBoundSamplerConfig(entropy_bound=entropy_bound)
|
|
sampler = EntropyBoundSampler(
|
|
config=sampler_config,
|
|
canvas_length=256,
|
|
vocab_size=10000,
|
|
max_denoising_steps=48,
|
|
)
|
|
return sampler
|