infra: changes to template

[sdhossain]

Changes

Main change is that I added a registry for attacks, and evals so that we can add an attack by doing the following:

@register_attack(AttackName.SOME_ATTACK)
Class SomeAttack(TamperAttack):
    """This is an attack implemented in SafeTuneBed"""
    ... rest of implementation

^ and the same pattern for evals.

This would allow this to be used as a package so that people could pip install tamperbench and then perhaps register an attack and run experiments on their end.

It also makes it a bit cleaner to do benchmarking later and add things, rather than having to populate dirs.

Testing

ICLR runs.

[sdhossain]

Note: some of these imports don't exist yet because I've retroactively split it up.

[tomtseng]

More simply I guess this could be
data["model_config"] = ModelConfig.from_dict(data["model_config"])? Also it looks like this modifies data in place — might not be desirable, should at least be documented as a side effect but might be better to do a deep clone

[tomtseng]

import shutil at top of file? it shouldn't be an expensive import

[tomtseng]

why re-create the directory after deleting?

[tomtseng]

Not sure H is necessary, it's not re-used several times like T

Suggested change

	name: AttackName, config_cls: type[H]
	name: AttackName, config_cls: TamperAttackConfig

[tomtseng]

as mentioned in the jailbreak tuning PR: for Qwen3, may want to also turn off thinking with <think> </think>, but does not apply to older Qwen models.

Print the output of

tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False  # Turns off thinking by adding <think></think>
)

to see what it looks like

[tomtseng]

pedantic preference but I'd find one of these easier to read:

TemplateName(name) if isinstance(name, str) else name

name if isinstance(name, TemplateName) else TemplateName(name)

[tomtseng]

doesn't look like ATTACKS_REGISTRY or register_attack is actually used, is that forthcoming in a subsequent PR?

[sdhossain]

Yeah it'd be used in forthcoming PRs where we do sweeps.

[tomtseng]

I kinda wonder if the registry pattern is useful/necessary for this case. Like, it looks equally valid to just define everything in-line so we're not doing this lazy evaluation thing where we define a function and then immediately call it to populate the template registry. (This feels a bit different from the attack/eval registries because here we're already just populating this registry here in this file, unlike attacks/evals where we're defining each attack & eval in separate files)

  _TEMPLATE_REGISTRY: dict[TemplateName, TextTemplate] = {
      TemplateName.LLAMA3: TextTemplate(
          user_prefix="<|start_header_id|>user<|end_header_id|>\n\n",
          assistant_prefix="<|start_header_id|>assistant<|end_header_id|>\n\n",
          end_turn="<|eot_id|>",
      ),
      # ... more templates inline
  }

  def get_template(name: str | TemplateName) -> TextTemplate:
      # ... same lookup logic

We could still have a function register_template(name: TemplateName, template: TextTemplate) doing _TEMPLATE_REGISTRY[name] = template if users want to add new templates at runtime.

[sdhossain]

It kind of makes it easier when using configs (i.e. you can specify in the .yaml file that we use a specific template).

[tomtseng]

Hmm so my understanding is that as long as _TEMPLATE_REGISTRY is populated in some way, then get_template() should be able to fetch specific templates, and I'm suggesting that _TEMPLATE_REGISTRY can be populated in a more direct way rather than making decorators do it.

So I guess what I'm actually suggesting is not to get rid of the registry pattern but rather to consider skipping using the decorator pattern.

[tomtseng]

similarly to attacks, I don't see this registry or @register_evaluation used anywhere, is it going to be a future PR?

[sdhossain]

Yeah I believe @register_evaluation will be added over in #35 and more specificially in: here

[sdhossain]

@register_evaluation(EvalName.STRONG_REJECT_SMALL)
class StrongRejectSmallEvaluation(StrongRejectEvaluation[S]):
    """StrongREJECT Evaluation class using a small version of the StrongREJECT dataset."""

    name: EvalName = EvalName.STRONG_REJECT_SMALL

    @override
    def load_strong_reject_prompts(self) -> list[str]:
        """Load the small version of the StrongReject dataset into an Arrow Dataset, and then return prompts.

        Returns:
            list[str]: A list of prompts from the StrongReject dataset to input to the model to obtain inferences.
        """
        strong_reject_dataset: ArrowDataset = (
            load_strong_reject_datasets.load_strongreject_small()
        )
        ...

and then we can specify the "strongreject_small" in a .yaml file to refer to the usage of this.

[tomtseng]

nit: why does only this import have noqa: F401?

[tomtseng]

nit: I wonder if we can also make the decorator set attack_cls.name:

@register_attack(AttackName.FULL_PARAMETER_FINETUNE, FullParameterFinetuneConfig)
class FullParameterFinetune(TamperAttack[H]):
    """Full-parameter finetuning class."""
    name: AttackName = AttackName.FULL_PARAMETER_FINETUNE  # if the decorator can set the name then could remove this line where the user has to specify the name yet again

(doesn't have to be this PR, since I know making changes with stacked PRs can end up being confusing)

[tomtseng]

I wonder if we could make use of the registry somehow instead of having to manually update stuff over here as well whenever we add a new eval. So like, instead of having 4 if statements, we'd have something like,

for eval_name, run_eval in EVALUATION_REGISTRY.items():
	if eval_name in self.attack_config.evals:
		results = pl.concat([results, run_eval(
    		model_checkpoint=self.output_checkpoint_path,
			...
    	)
    ])

[sdhossain]

Closing as we merged in #53