e412370877
* roundup_power2_divisions not needed with newer pytorch versions * remove typo * update qwen3.5 moe 35b-a3b yaml for 5090 * more bug fixes * fix tests to match updated trainer * don't use fa2 for hooks test * reset plugins on the instance * retry download * fix references to renamed axolotl_cfg property on trainer * Fix ref to trainer cfg
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
"""Tests for diffusion generation callback dataloader selection and triggering."""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from axolotl.integrations.diffusion import DiffusionGenerationCallback
|
|
|
|
|
|
class DummyTrainer:
|
|
"""Minimal trainer double with required attributes/methods for the callback."""
|
|
|
|
def __init__(self, use_eval: bool):
|
|
# Config used by callback
|
|
self.axolotl_cfg = SimpleNamespace(
|
|
diffusion=SimpleNamespace(
|
|
generation_interval=1,
|
|
num_generation_samples=1,
|
|
generation_max_length=32,
|
|
generation_steps=4,
|
|
generation_temperature=0.0,
|
|
mask_token_id=16,
|
|
),
|
|
use_wandb=False,
|
|
)
|
|
|
|
# Model/tokenizer are passed through to generate_samples; not used here
|
|
self.model = Mock()
|
|
self.processing_class = Mock()
|
|
|
|
# Datasets and loaders
|
|
self.eval_dataset = object() if use_eval else None
|
|
self._train_loader = object()
|
|
self._eval_loader = object()
|
|
|
|
# State for world process check
|
|
self.state = SimpleNamespace(is_world_process_zero=True)
|
|
|
|
# Track which loader was requested
|
|
self.requested: list[str] = []
|
|
|
|
def get_train_dataloader(self):
|
|
self.requested.append("train")
|
|
return self._train_loader
|
|
|
|
def get_eval_dataloader(self):
|
|
self.requested.append("eval")
|
|
return self._eval_loader
|
|
|
|
|
|
@pytest.mark.parametrize("use_eval", [False, True])
|
|
def test_callback_uses_correct_dataloader(monkeypatch, use_eval):
|
|
trainer = DummyTrainer(use_eval=use_eval)
|
|
callback = DiffusionGenerationCallback(trainer)
|
|
|
|
captured = {}
|
|
|
|
# Patch generate_samples in the callback module's namespace
|
|
def fake_generate_samples(**kwargs):
|
|
captured["dataloader"] = kwargs.get("dataloader")
|
|
# Return one dummy sample to exercise logging path
|
|
return [
|
|
{
|
|
"original": "o",
|
|
"masked": "m",
|
|
"generated": "g",
|
|
"mask_ratio": 0.5,
|
|
"masked_tokens": 1,
|
|
"total_tokens": 2,
|
|
}
|
|
]
|
|
|
|
monkeypatch.setattr(
|
|
"axolotl.integrations.diffusion.callbacks.generate_samples",
|
|
fake_generate_samples,
|
|
)
|
|
|
|
# Trigger at step 1 (interval=1)
|
|
args = SimpleNamespace()
|
|
state = SimpleNamespace(global_step=1)
|
|
control = SimpleNamespace()
|
|
|
|
callback.on_step_end(args=args, state=state, control=control)
|
|
|
|
# Assert the expected dataloader path was used
|
|
if use_eval:
|
|
assert trainer.requested[0] == "eval"
|
|
assert captured["dataloader"] is trainer._eval_loader
|
|
else:
|
|
assert trainer.requested[0] == "train"
|
|
assert captured["dataloader"] is trainer._train_loader
|