Gradient-Toxicity-Detection/score_data.py at master · olo126/Gradient-Toxicity-Detection · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import argparse
import os

import torch

argparser = argparse.ArgumentParser(
    description='Script for selecting the data for training')
argparser.add_argument('--gradient_path', type=str, default="{} ckpt{}",
                       help='The path to the gradient file')
argparser.add_argument('--train_file_names', type=str, nargs='+',
                       help='The name of the training file')
argparser.add_argument('--ckpts', type=int, nargs='+',
                       help="Checkpoint numbers.")
argparser.add_argument('--checkpoint_weights', type=float, nargs='+',
                       help="checkpoint weights")
argparser.add_argument('--target_task_names', type=str,
                       nargs='+', help="The name of the target tasks")
argparser.add_argument('--validation_gradient_path', type=str,
                       default="{} ckpt{}", help='The path to the validation gradient file')
argparser.add_argument('--output_path', type=str, default="selected_data",
                       help='The path to the output')
argparser.add_argument('--separate_scores', action="store_true",
                       help='Separate scores for each validation point')
argparser.add_argument('--combine_influence', action="store_true",
                       help='Combines toxic and benign influence scores')


args = argparser.parse_args()

N_SUBTASKS = {"toxic": 1, "benign": 1}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def calculate_influence_score(training_info: torch.Tensor, validation_info: torch.Tensor):
    """Calculate the influence score.

    Args:
        training_info (torch.Tensor): training info (gradients/representations) stored in a tensor of shape N x N_DIM
        validation_info (torch.Tensor): validation info (gradients/representations) stored in a tensor of shape N_VALID x N_DIM
    """
    # N x N_VALID
    influence_scores = torch.matmul(
        training_info, validation_info.transpose(0, 1))
    return influence_scores

# renormalize the checkpoint weights
if sum(args.checkpoint_weights) != 1:
    s = sum(args.checkpoint_weights)
    args.checkpoint_weights = [i/s for i in args.checkpoint_weights]

# calculate the influence score for each validation task
for target_task_name in args.target_task_names:
    if args.combine_influence:
        for train_file_name in args.train_file_names:
            influence_score = 0
            for i, ckpt in enumerate(args.ckpts):
                #first set of influence (toxic)
                validation_path = args.validation_gradient_path.format(
                    target_task_name, ckpt)
                if os.path.isdir(validation_path):
                    validation_path = os.path.join(validation_path, "all_orig.pt")
                validation_info = torch.load(validation_path)

                if not torch.is_tensor(validation_info):
                    validation_info = torch.tensor(validation_info)
                validation_info = validation_info.to(device).float()
                gradient_path = args.gradient_path.format(train_file_name, ckpt)
                if os.path.isdir(gradient_path):
                    gradient_path = os.path.join(gradient_path, "all_orig.pt")
                training_info = torch.load(gradient_path)

                if not torch.is_tensor(training_info):
                    training_info = torch.tensor(training_info)
                training_info = training_info.to(device).float()

                #second set of influence (benign)
                validation_path_2 = args.validation_gradient_path.format(
                    args.target_task_names[-1], ckpt)
                if os.path.isdir(validation_path_2):
                    validation_path_2 = os.path.join(validation_path_2, "all_orig.pt")
                validation_info_2 = torch.load(validation_path_2)

                if not torch.is_tensor(validation_info_2):
                    validation_info_2 = torch.tensor(validation_info_2)
                validation_info_2 = validation_info_2.to(device).float()

                combined_influence = (args.checkpoint_weights[i] * \
                    calculate_influence_score(
                        training_info=training_info, validation_info=validation_info)) - \
                    (0.5 * args.checkpoint_weights[i] * \
                    calculate_influence_score(
                        training_info=training_info, validation_info=validation_info_2))
                influence_score += combined_influence
            print(influence_score.shape)
            if not args.separate_scores:
                print("combining scores")
                influence_score = influence_score.reshape(
                    influence_score.shape[0], N_SUBTASKS[target_task_name], -1).mean(-1).max(-1)[0]
                output_dir = os.path.join(args.output_path, target_task_name)
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                output_file = os.path.join(
                    args.output_path, target_task_name, f"{train_file_name}_influence_score.pt")
                torch.save(influence_score, output_file)
                print("Saved influence score to {}".format(output_file))
            else:
                for i in range(influence_score.shape[1]):
                    target_task_name_sep = os.path.join(target_task_name, f"val_{i}")
                    output_dir = os.path.join(args.output_path, target_task_name_sep)
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    output_file = os.path.join(
                        args.output_path, target_task_name_sep, f"{train_file_name}_influence_score_val_{i}.pt")
                    torch.save(influence_score[:,i], output_file)
                    print("Saved influence score to {}".format(output_file))
        break
    else:
        for train_file_name in args.train_file_names:
            influence_score = 0
            for i, ckpt in enumerate(args.ckpts):
                # validation_path = args.validation_gradient_path.format(
                # target_task_name, ckpt)
                validation_path = args.validation_gradient_path.format(
                    target_task_name, ckpt)
                if os.path.isdir(validation_path):
                    validation_path = os.path.join(validation_path, "all_orig.pt")
                validation_info = torch.load(validation_path)

                if not torch.is_tensor(validation_info):
                    validation_info = torch.tensor(validation_info)
                validation_info = validation_info.to(device).float()
                # gradient_path = args.gradient_path.format(train_file_name, ckpt)
                gradient_path = args.gradient_path.format(train_file_name, ckpt)
                if os.path.isdir(gradient_path):
                    gradient_path = os.path.join(gradient_path, "all_orig.pt")
                training_info = torch.load(gradient_path)

                if not torch.is_tensor(training_info):
                    training_info = torch.tensor(training_info)
                training_info = training_info.to(device).float()

                influence_score += args.checkpoint_weights[i] * \
                    calculate_influence_score(
                        training_info=training_info, validation_info=validation_info)
            print(influence_score.shape)
            if not args.separate_scores:
                print("combining scores")
                influence_score = influence_score.reshape(
                    influence_score.shape[0], N_SUBTASKS[target_task_name], -1).mean(-1).max(-1)[0]
                output_dir = os.path.join(args.output_path, target_task_name)
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                output_file = os.path.join(
                    args.output_path, target_task_name, f"{train_file_name}_influence_score.pt")
                torch.save(influence_score, output_file)
                print("Saved influence score to {}".format(output_file))
            else:
                for i in range(influence_score.shape[1]):
                    target_task_name_sep = os.path.join(target_task_name, f"val_{i}")
                    output_dir = os.path.join(args.output_path, target_task_name_sep)
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    output_file = os.path.join(
                        args.output_path, target_task_name_sep, f"{train_file_name}_influence_score_val_{i}.pt")
                    torch.save(influence_score[:,i], output_file)
                    print("Saved influence score to {}".format(output_file))