diff --git a/src/soup_cli/trainer/grpo.py b/src/soup_cli/trainer/grpo.py index e185b47..751a984 100644 --- a/src/soup_cli/trainer/grpo.py +++ b/src/soup_cli/trainer/grpo.py @@ -489,7 +489,13 @@ class GRPOTrainerWrapper: model_kwargs["quantization_config"] = quant_config_obj self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs) + from soup_cli.utils.data_pipeline import apply_vocab_expansion + apply_vocab_expansion( + self.tokenizer, + self.model, + cfg.data, + ) if tcfg.quantization in ("4bit", "8bit", "mxfp4"): self.model = prepare_model_for_kbit_training(self.model) diff --git a/src/soup_cli/trainer/orpo.py b/src/soup_cli/trainer/orpo.py index a05ad91..47a3154 100644 --- a/src/soup_cli/trainer/orpo.py +++ b/src/soup_cli/trainer/orpo.py @@ -197,7 +197,13 @@ class ORPOTrainerWrapper: model_kwargs["quantization_config"] = quant_config_obj self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs) + from soup_cli.utils.data_pipeline import apply_vocab_expansion + apply_vocab_expansion( + self.tokenizer, + self.model, + cfg.data, + ) if tcfg.quantization in ("4bit", "8bit", "mxfp4"): self.model = prepare_model_for_kbit_training(self.model) diff --git a/src/soup_cli/trainer/simpo.py b/src/soup_cli/trainer/simpo.py index d50143c..35f3991 100644 --- a/src/soup_cli/trainer/simpo.py +++ b/src/soup_cli/trainer/simpo.py @@ -200,7 +200,13 @@ class SimPOTrainerWrapper: model_kwargs["quantization_config"] = quant_config_obj self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs) + from soup_cli.utils.data_pipeline import apply_vocab_expansion + apply_vocab_expansion( + self.tokenizer, + self.model, + cfg.data, + ) if tcfg.quantization in ("4bit", "8bit", "mxfp4"): self.model = prepare_model_for_kbit_training(self.model) diff --git a/tests/test_trainer_init.py b/tests/test_trainer_init.py index fbe8a77..429615d 100644 --- a/tests/test_trainer_init.py +++ b/tests/test_trainer_init.py @@ -499,6 +499,114 @@ class TestGRPOTrainerInit: assert wrapper.model is None assert wrapper.tokenizer is None + def test_grpo_vocab_expansion_adds_tokens_and_resizes(self, monkeypatch): + """GRPO should apply configured vocabulary expansion.""" + + import sys + import types + from types import SimpleNamespace + + from soup_cli.trainer.grpo import GRPOTrainerWrapper + + cfg = _make_config( + task="grpo", + data={ + "train": "./data.jsonl", + "format": "chatml", + "add_new_tokens": ["", ""], + "new_special_tokens": [""], + "resize_vocab": True, + }, + training={"quantization": "none"}, + ) + + calls = { + "add_tokens": None, + "add_special_tokens": None, + "resize": None, + } + + class _Tokenizer: + pad_token = None + eos_token = "" + + def __init__(self): + self.vocab = {"": 0} + + def get_vocab(self): + return self.vocab + + def add_tokens(self, tokens): + calls["add_tokens"] = list(tokens) + added = 0 + for token in tokens: + if token not in self.vocab: + self.vocab[token] = len(self.vocab) + added += 1 + return added + + def add_special_tokens(self, data): + tokens = list(data["additional_special_tokens"]) + calls["add_special_tokens"] = tokens + added = 0 + for token in tokens: + if token not in self.vocab: + self.vocab[token] = len(self.vocab) + added += 1 + return added + + def __len__(self): + return len(self.vocab) + + class _Model: + config = SimpleNamespace() + + def resize_token_embeddings(self, size): + calls["resize"] = size + + def parameters(self): + return [] + + tokenizer = _Tokenizer() + model = _Model() + + fake_transformers = types.SimpleNamespace( + AutoTokenizer=types.SimpleNamespace( + from_pretrained=lambda *args, **kwargs: tokenizer + ), + AutoModelForCausalLM=types.SimpleNamespace( + from_pretrained=lambda *args, **kwargs: model + ), + ) + + fake_peft = types.SimpleNamespace( + LoraConfig=lambda **kwargs: SimpleNamespace(**kwargs), + TaskType=SimpleNamespace(CAUSAL_LM="CAUSAL_LM"), + get_peft_model=lambda model_obj, _cfg: model_obj, + prepare_model_for_kbit_training=lambda model_obj: model_obj, + ) + + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + monkeypatch.setitem(sys.modules, "peft", fake_peft) + + monkeypatch.setattr( + "soup_cli.utils.quant_menu.build_quantization_config_for_loader", + lambda **kwargs: None, + ) + + wrapper = object.__new__(GRPOTrainerWrapper) + wrapper.config = cfg + wrapper.device = "cpu" + wrapper._trust_remote_code = False + wrapper.model = None + wrapper.tokenizer = None + + wrapper._setup_transformers(cfg, cfg.training) + + assert calls["add_tokens"] == [""] + assert calls["add_special_tokens"] == [""] + assert calls["resize"] == len(tokenizer) + class TestRewardModelTrainerInit: """Test RewardModelTrainerWrapper constructor.""" @@ -853,6 +961,224 @@ class TestBCOTrainerInit: assert calls["add_special_tokens"] == [""] assert calls["resize"] == len(tokenizer) +class TestORPOTrainerInit: + def test_orpo_vocab_expansion_adds_tokens_and_resizes(self, monkeypatch): + """ORPO should apply configured vocabulary expansion.""" + + import sys + import types + from types import SimpleNamespace + + from soup_cli.trainer.orpo import ORPOTrainerWrapper + + cfg = _make_config( + task="orpo", + data={ + "train": "./data.jsonl", + "format": "chatml", + "add_new_tokens": ["", ""], + "new_special_tokens": [""], + "resize_vocab": True, + }, + training={"quantization": "none"}, + ) + + calls = { + "add_tokens": None, + "add_special_tokens": None, + "resize": None, + } + + class _Tokenizer: + pad_token = None + eos_token = "" + + def __init__(self): + self.vocab = {"": 0} + + def get_vocab(self): + return self.vocab + + def add_tokens(self, tokens): + calls["add_tokens"] = list(tokens) + added = 0 + for token in tokens: + if token not in self.vocab: + self.vocab[token] = len(self.vocab) + added += 1 + return added + + def add_special_tokens(self, data): + tokens = list(data["additional_special_tokens"]) + calls["add_special_tokens"] = tokens + added = 0 + for token in tokens: + if token not in self.vocab: + self.vocab[token] = len(self.vocab) + added += 1 + return added + + def __len__(self): + return len(self.vocab) + + class _Model: + config = SimpleNamespace() + + def resize_token_embeddings(self, size): + calls["resize"] = size + + def parameters(self): + return [] + + tokenizer = _Tokenizer() + model = _Model() + + fake_transformers = types.SimpleNamespace( + AutoTokenizer=types.SimpleNamespace( + from_pretrained=lambda *args, **kwargs: tokenizer + ), + AutoModelForCausalLM=types.SimpleNamespace( + from_pretrained=lambda *args, **kwargs: model + ), + ) + + fake_peft = types.SimpleNamespace( + LoraConfig=lambda **kwargs: SimpleNamespace(**kwargs), + TaskType=SimpleNamespace(CAUSAL_LM="CAUSAL_LM"), + get_peft_model=lambda model_obj, _cfg: model_obj, + prepare_model_for_kbit_training=lambda model_obj: model_obj, + ) + + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + monkeypatch.setitem(sys.modules, "peft", fake_peft) + + monkeypatch.setattr( + "soup_cli.utils.quant_menu.build_quantization_config_for_loader", + lambda **kwargs: None, + ) + + wrapper = object.__new__(ORPOTrainerWrapper) + wrapper.config = cfg + wrapper.device = "cpu" + wrapper._trust_remote_code = False + wrapper.model = None + wrapper.tokenizer = None + + wrapper._setup_transformers(cfg, cfg.training) + + assert calls["add_tokens"] == [""] + assert calls["add_special_tokens"] == [""] + assert calls["resize"] == len(tokenizer) + +class TestSIMPOTrainerInit: + def test_simpo_vocab_expansion_adds_tokens_and_resizes(self, monkeypatch): + """SIMPO should apply configured vocabulary expansion.""" + + import sys + import types + from types import SimpleNamespace + + from soup_cli.trainer.simpo import SimPOTrainerWrapper + + cfg = _make_config( + task="simpo", + data={ + "train": "./data.jsonl", + "format": "chatml", + "add_new_tokens": ["", ""], + "new_special_tokens": [""], + "resize_vocab": True, + }, + training={"quantization": "none"}, + ) + + calls = { + "add_tokens": None, + "add_special_tokens": None, + "resize": None, + } + + class _Tokenizer: + pad_token = None + eos_token = "" + + def __init__(self): + self.vocab = {"": 0} + + def get_vocab(self): + return self.vocab + + def add_tokens(self, tokens): + calls["add_tokens"] = list(tokens) + added = 0 + for token in tokens: + if token not in self.vocab: + self.vocab[token] = len(self.vocab) + added += 1 + return added + + def add_special_tokens(self, data): + tokens = list(data["additional_special_tokens"]) + calls["add_special_tokens"] = tokens + added = 0 + for token in tokens: + if token not in self.vocab: + self.vocab[token] = len(self.vocab) + added += 1 + return added + + def __len__(self): + return len(self.vocab) + + class _Model: + config = SimpleNamespace() + + def resize_token_embeddings(self, size): + calls["resize"] = size + + def parameters(self): + return [] + + tokenizer = _Tokenizer() + model = _Model() + + fake_transformers = types.SimpleNamespace( + AutoTokenizer=types.SimpleNamespace( + from_pretrained=lambda *args, **kwargs: tokenizer + ), + AutoModelForCausalLM=types.SimpleNamespace( + from_pretrained=lambda *args, **kwargs: model + ), + ) + + fake_peft = types.SimpleNamespace( + LoraConfig=lambda **kwargs: SimpleNamespace(**kwargs), + TaskType=SimpleNamespace(CAUSAL_LM="CAUSAL_LM"), + get_peft_model=lambda model_obj, _cfg: model_obj, + prepare_model_for_kbit_training=lambda model_obj: model_obj, + ) + + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + monkeypatch.setitem(sys.modules, "peft", fake_peft) + + monkeypatch.setattr( + "soup_cli.utils.quant_menu.build_quantization_config_for_loader", + lambda **kwargs: None, + ) + + wrapper = object.__new__(SimPOTrainerWrapper) + wrapper.config = cfg + wrapper.device = "cpu" + wrapper._trust_remote_code = False + wrapper.model = None + wrapper.tokenizer = None + + wrapper._setup_transformers(cfg, cfg.training) + + assert calls["add_tokens"] == [""] + assert calls["add_special_tokens"] == [""] + assert calls["resize"] == len(tokenizer) + class TestTrainTaskRouting: """Test that train command routes to correct trainer based on task."""