def log()

in pyrit/auxiliary_attacks/gcg/attack/base/attack_manager.py [0:0]

• 56 lines of code
• 9 McCabe index (conditional complexity)

    def log(self, step_num, n_steps, control, loss, runtime, model_tests, verbose=True):

        prompt_tests_jb, prompt_tests_mb, model_tests_loss = list(map(np.array, model_tests))
        all_goal_strs = self.goals + self.test_goals
        all_workers = self.workers + self.test_workers
        tests = {
            all_goal_strs[i]: [
                (
                    all_workers[j].model.name_or_path,
                    prompt_tests_jb[j][i],
                    prompt_tests_mb[j][i],
                    model_tests_loss[j][i],
                )
                for j in range(len(all_workers))
            ]
            for i in range(len(all_goal_strs))
        }
        n_passed = self.parse_results(prompt_tests_jb)
        n_em = self.parse_results(prompt_tests_mb)
        n_loss = self.parse_results(model_tests_loss)
        total_tests = self.parse_results(np.ones(prompt_tests_jb.shape, dtype=int))
        n_loss = [lo / t if t > 0 else 0 for lo, t in zip(n_loss, total_tests)]

        tests["n_passed"] = n_passed
        tests["n_em"] = n_em
        tests["n_loss"] = n_loss
        tests["total"] = total_tests

        with open(self.logfile, "r") as f:
            log = json.load(f)

        log["controls"].append(control)
        log["losses"].append(loss)
        log["runtimes"].append(runtime)
        log["tests"].append(tests)

        with open(self.logfile, "w") as f:
            json.dump(log, f, indent=4, cls=NpEncoder)

        if verbose:
            output_str = ""
            for i, tag in enumerate(["id_id", "id_od", "od_id", "od_od"]):
                if total_tests[i] > 0:
                    output_str += (
                        f"({tag}) | Passed {n_passed[i]:>3}/{total_tests[i]:<3} | "
                        f"EM {n_em[i]:>3}/{total_tests[i]:<3} | "
                        f"Loss {n_loss[i]:.4f}\n"
                    )
            print(
                (
                    f"\n====================================================\n"
                    f"Step {step_num:>4}/{n_steps:>4} ({runtime:.4} s)\n"
                    f"{output_str}"
                    f"control='{control}'\n"
                    f"====================================================\n"
                )
            )

        # Log to mlflow
        log_loss(step=step_num, loss=loss)
        log_gpu_memory(step=step_num)

        # Log results table to mlflow
        if step_num == n_steps:
            log_table_summary(losses=log["losses"], controls=log["controls"], n_steps=n_steps)
            mlflow.end_run()