test.py

import torch
import hydra
from omegaconf import DictConfig
import numpy as np
import random
from src.trainer import Trainer
import wandb
from omegaconf import OmegaConf
import os
from dotenv import load_dotenv

VALIDATE_AFTER = 10
INDENT = ' ' * 5

@hydra.main(config_path="configs", config_name="config_train.yaml", version_base="1.1")
def main(config: DictConfig) -> None:
    load_dotenv()
    print (config)
    
    # os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(lambda x: x.split(':')[1], config.device_ids))
    
    wandb.config = OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
    # OmegaConf.register_new_resolver("eval", eval)

    run = wandb.init(
        project=config.wandb.project,
        entity=config.wandb.entity,
        group=config.group,
        mode=config.wandb.mode if 'mode' in config.wandb else None, #"disabled",
        config=wandb.config
    )

    save_dir = f"{config.log_dir}/saved_models/{config.language_model.lm_name}_{config.defense['_target_'].split('.')[-1]}/"
    save_dir = f'{save_dir}/{config.save_dir}/' if 'save_dir' in config else f'{save_dir}'

    
    trainer = Trainer(language_model=config.language_model,
                      attack=config.attack,
                      defense=config.defense,
                      dataset=config.dataset,
                      train_cfg=config.train_params,
                      save_dir=save_dir,
                      seed=config.seed)
    
    trainer.load_defense()
    
    trainer.evaluate_all()
    
    wandb.finish()
    
    
if __name__ == "__main__":
    main()