Spaces:
Paused
Paused
| # 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. | |
| r"""NPO Example. | |
| NPO [https://arxiv.org/pdf/2404.05868] is an unlearning algorithm that is | |
| a modification of DPO. Like DPO, it uses two models to generate the response. | |
| But NPO only focuses on reducing the answer likelihood compared to the reference model. | |
| L_NPO = − E[log σ(− β log{π_θ(y|x)/π_ref (y|x)}) | |
| Train locally with: | |
| ```sh | |
| python -m kauldron.main \ | |
| --cfg=examples/npo.py \ | |
| --cfg.workdir=/tmp/kauldron_oss/workdir | |
| ``` | |
| """ | |
| from kauldron import konfig | |
| # pylint: disable=g-import-not-at-top | |
| with konfig.imports(): | |
| from gemma import gm | |
| from kauldron import kd | |
| import optax | |
| import tensorflow_datasets as tfds | |
| # pylint: enable=g-import-not-at-top | |
| def get_config(): | |
| """Get the default hyperparameter configuration.""" | |
| return kd.train.Trainer( | |
| seed=42, | |
| # Dataset | |
| train_ds=_make_dataset( | |
| dataset_name="locuslab__tofu/forget10", | |
| data_dir=tfds.HF_PLACER_DIR, | |
| ), | |
| # Model definition | |
| model=gm.nn.AnchoredPolicy( | |
| policy=gm.nn.Gemma3_12B(tokens="batch.input", text_only=True), | |
| ), | |
| # Load the weights from the pretrained checkpoint | |
| init_transform=gm.ckpts.AnchoredPolicyLoader( | |
| policy=gm.ckpts.LoadCheckpoint( | |
| path=gm.ckpts.CheckpointPath.GEMMA3_12B_IT, | |
| ), | |
| ), | |
| # Training | |
| num_train_steps=250, | |
| train_losses={ | |
| "npo": gm.losses.NpoLoss( | |
| tokens="batch.target", | |
| sequence_mask="batch.mask", | |
| policy_logits="preds.policy.logits", | |
| anchor_logits="preds.anchor.logits", | |
| ), | |
| }, | |
| optimizer=optax.adafactor(learning_rate=1e-4), | |
| checkpointer=kd.ckpts.Checkpointer( | |
| save_interval_steps=50, | |
| ), | |
| # Evaluation | |
| evals={ | |
| # "test": kd.evals.Evaluator( | |
| # run=kd.evals.EveryNSteps(1000), | |
| # ds=_make_dataset( | |
| # training=False, dataset_name=dataset_name, data_dir=data_dir | |
| # ), | |
| # ), | |
| }, | |
| ) | |
| def _make_dataset( | |
| dataset_name: str, data_dir: str | |
| ): | |
| """Make a dataset for training or evaluation. | |
| Args: | |
| dataset_name: The name of the dataset to load. | |
| data_dir: The directory where the dataset is located. | |
| Returns: | |
| A `kd.data.Pipeline` for the dataset. | |
| """ | |
| max_length = 256 | |
| batch_size = 16 | |
| tokenizer = gm.text.Gemma3Tokenizer() | |
| return kd.data.py.Tfds( | |
| name=dataset_name, | |
| data_dir=data_dir, | |
| split="all", | |
| shuffle=True, | |
| batch_size=batch_size, | |
| transforms=[ | |
| # Drop the 'split' column | |
| kd.data.py.Elements(keep=["question", "answer"]), | |
| # Create the model inputs/targets/mask. | |
| gm.data.Seq2SeqTask( | |
| in_prompt="question", | |
| in_response="answer", | |
| # Output batch is {'input': ..., 'target': ..., 'mask': ...} | |
| out_input="input", | |
| out_target="target", | |
| out_target_mask="mask", | |
| tokenizer=tokenizer, | |
| # Padding parameters | |
| max_length=max_length, | |
| truncate=True, | |
| ), | |
| ], | |
| ) | |