Skip to content

Commit a9a8f0f

Browse files
madlagmadlag
committed
Work on command line utility.
1 parent 7776d91 commit a9a8f0f

1 file changed

Lines changed: 101 additions & 99 deletions

File tree

‎nn_pruning/command_line.py‎

Lines changed: 101 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -1,113 +1,121 @@
11
import click
22
import json
3-
from pathlib import Path
3+
44

55
@click.group()
66
@click.pass_context
77
def cli(ctx):
88
ctx.obj = {}
99

1010

11-
1211
QA_TYPICAL_PARAMETERS = {
13-
"dataset_name": "squad",
14-
"do_train": 1,
15-
"do_eval": 1,
16-
17-
"max_seq_length": 384,
18-
"doc_stride": 128,
19-
"num_train_epochs": 20,
20-
"logging_steps": 250,
21-
"save_steps": 10000,
22-
"eval_steps": 10000,
23-
"save_total_limit": 50,
24-
"seed": 17,
25-
"evaluation_strategy": "steps",
26-
"learning_rate": 3e-5,
27-
"mask_scores_learning_rate": 1e-2,
28-
"overwrite_cache": 0,
29-
"overwrite_output_dir": 1,
30-
"warmup_steps": 5400,
31-
"initial_warmup": 1,
32-
"final_warmup": 10,
33-
"initial_threshold": 0,
34-
"final_threshold": 0.1,
35-
"dense_pruning_method": "sigmoied_threshold:1d_alt",
36-
"dense_block_rows":1,
37-
"dense_block_cols":1,
38-
39-
"attention_pruning_method": "sigmoied_threshold",
40-
"attention_block_rows":32,
41-
"attention_block_cols":32,
42-
"attention_lambda":1.0,
43-
44-
"mask_init": "constant",
45-
"mask_scale": 0.0,
46-
"regularization": "l1",
47-
"distil_alpha_ce": 0.1,
48-
"distil_alpha_teacher": 0.9,
49-
"attention_output_with_dense": 0
12+
"dataset_name": "squad",
13+
"do_train": 1,
14+
"do_eval": 1,
15+
"max_seq_length": 384,
16+
"doc_stride": 128,
17+
"num_train_epochs": 20,
18+
"logging_steps": 250,
19+
"save_steps": 10000,
20+
"eval_steps": 10000,
21+
"save_total_limit": 50,
22+
"seed": 17,
23+
"evaluation_strategy": "steps",
24+
"learning_rate": 3e-5,
25+
"mask_scores_learning_rate": 1e-2,
26+
"overwrite_cache": 0,
27+
"overwrite_output_dir": 1,
28+
"warmup_steps": 5400,
29+
"initial_warmup": 1,
30+
"final_warmup": 10,
31+
"initial_threshold": 0,
32+
"final_threshold": 0.1,
33+
"dense_pruning_method": "sigmoied_threshold:1d_alt",
34+
"dense_block_rows": 1,
35+
"dense_block_cols": 1,
36+
"attention_pruning_method": "sigmoied_threshold",
37+
"attention_block_rows": 32,
38+
"attention_block_cols": 32,
39+
"attention_lambda": 1.0,
40+
"mask_init": "constant",
41+
"mask_scale": 0.0,
42+
"regularization": "l1",
43+
"distil_alpha_ce": 0.1,
44+
"distil_alpha_teacher": 0.9,
45+
"attention_output_with_dense": 0,
5046
}
5147

5248
GLUE_TYPICAL_PARAMETERS = {
53-
"task_name": "mnli",
54-
"do_train": 1,
55-
"do_eval": 1,
56-
"per_device_eval_batch_size": 128,
57-
"max_seq_length": 128,
58-
"doc_stride": 128,
59-
"num_train_epochs": 12,
60-
"logging_steps": 250,
61-
"save_steps": 5000,
62-
"eval_steps": 5000,
63-
"save_total_limit": 50,
64-
"seed": 17,
65-
"evaluation_strategy": "steps",
66-
"learning_rate": 3e-5,
67-
"mask_scores_learning_rate": 1e-2,
68-
"overwrite_cache": 0,
69-
"overwrite_output_dir": 1,
70-
"warmup_steps": 12000,
71-
"initial_warmup": 1,
72-
"final_warmup": 4,
73-
"initial_threshold": 0,
74-
"final_threshold": 0.1,
75-
"dense_pruning_method": "sigmoied_threshold:1d_alt",
76-
"dense_block_rows":1,
77-
"dense_block_cols":1,
78-
"attention_pruning_method": "sigmoied_threshold",
79-
"attention_block_rows":32,
80-
"attention_block_cols":32,
81-
"attention_lambda":1.0,
82-
"mask_init": "constant",
83-
"mask_scale": 0.0,
84-
"regularization": "l1",
85-
"distil_alpha_ce": 0.1,
86-
"distil_alpha_teacher": 0.90,
87-
"attention_output_with_dense": 0
49+
"task_name": "mnli",
50+
"do_train": 1,
51+
"do_eval": 1,
52+
"per_device_eval_batch_size": 128,
53+
"max_seq_length": 128,
54+
"doc_stride": 128,
55+
"num_train_epochs": 12,
56+
"logging_steps": 250,
57+
"save_steps": 5000,
58+
"eval_steps": 5000,
59+
"save_total_limit": 50,
60+
"seed": 17,
61+
"evaluation_strategy": "steps",
62+
"learning_rate": 3e-5,
63+
"mask_scores_learning_rate": 1e-2,
64+
"overwrite_cache": 0,
65+
"overwrite_output_dir": 1,
66+
"warmup_steps": 12000,
67+
"initial_warmup": 1,
68+
"final_warmup": 4,
69+
"initial_threshold": 0,
70+
"final_threshold": 0.1,
71+
"dense_pruning_method": "sigmoied_threshold:1d_alt",
72+
"dense_block_rows": 1,
73+
"dense_block_cols": 1,
74+
"attention_pruning_method": "sigmoied_threshold",
75+
"attention_block_rows": 32,
76+
"attention_block_cols": 32,
77+
"attention_lambda": 1.0,
78+
"mask_init": "constant",
79+
"mask_scale": 0.0,
80+
"regularization": "l1",
81+
"distil_alpha_ce": 0.1,
82+
"distil_alpha_teacher": 0.90,
83+
"attention_output_with_dense": 0,
8884
}
8985

90-
#parameters: model_name, teacher, batch_size, regu_lambda, regu_lambda_dense, ampere
91-
86+
# parameters: model_name, teacher, batch_size, regu_lambda, regu_lambda_dense, ampere
87+
@cli.group()
88+
def examples():
89+
pass
9290

93-
94-
@cli.command()
91+
@examples.command()
9592
@click.pass_context
96-
@click.argument("task", default="squadv1", type=click.Choice(['squadv1', 'mnli']))
97-
@click.option("--json-path", type=click.Path(resolve_path=True))
98-
@click.option("--model-name-or-path", defaut="bert-base-uncased", type=click.Choice("bert-base-uncased", "bert-large-uncased"))
99-
@click.option("--distil-teacher-name-or-path-name", defaut=None, type=str)
100-
@click.option("--per-device-train-batch-size", defaut=16, type=int)
101-
@click.option("--regularization-final-lambda", defaut=20, type=float)
102-
@click.option("--dense-lambda", defaut=1.0, type=float)
103-
@click.option("--ampere-pruning-method", default="disabled", type=click.Choice("disabled", "topk"))
104-
@click.option("--output-dir", type=click.Path(resolve_path=True))
105-
def fine_prune(ctx, task, json_path, model_name_or_path, distil_teacher_name_or_path, per_device_train_batch_size, regularization_final_lambda, dense_lambda, ampere_pruning_method, output_dir):
93+
@click.argument("task", default="squadv1", type=click.Choice(["squadv1", "mnli"]))
94+
@click.argument("output-dir", type=click.Path(resolve_path=True))
95+
@click.option("--parameters", type=click.Path(resolve_path=True), help="Path to a parameters json file")
96+
@click.option("--model-name-or-path", default="bert-base-uncased", type=click.Choice(["bert-base-uncased", "bert-large-uncased"]))
97+
@click.option("--distil-teacher-name-or-path", default="bert-large-uncased-whole-word-masking-finetuned-squad", type=str, help = "teacher name (default is bert-large-uncased-whole-word-masking-finetuned-squad)")
98+
@click.option("--per-device-train-batch-size", default=16, type=int)
99+
@click.option("--regularization-final-lambda", default=10, type=float)
100+
@click.option("--dense-lambda", default=1.0, type=float)
101+
@click.option("--ampere-pruning-method", default="disabled", type=click.Choice(["disabled", "topk"]))
102+
def fine_prune(
103+
ctx,
104+
task,
105+
json_path,
106+
model_name_or_path,
107+
distil_teacher_name_or_path,
108+
per_device_train_batch_size,
109+
regularization_final_lambda,
110+
dense_lambda,
111+
ampere_pruning_method,
112+
output_dir,
113+
):
106114
filename = json_path
107-
if json_path is None:
115+
if json_path is not None:
108116
param_dict = json.load(open(filename))
109117
else:
110-
if task == "squad":
118+
if task == "squadv1":
111119
param_dict = QA_TYPICAL_PARAMETERS
112120
elif task == "mnli":
113121
param_dict = GLUE_TYPICAL_PARAMETERS
@@ -129,26 +137,20 @@ def fine_prune(ctx, task, json_path, model_name_or_path, distil_teacher_name_or_
129137
param_dict["distil_teacher_name_or_path"] = "csarron/bert-base-uncased-squad-v1"
130138

131139
import nn_pruning.examples.question_answering.qa_sparse_xp as qa_sparse_xp
140+
132141
experiment = qa_sparse_xp.QASparseXP(param_dict)
133142
else:
134143
if json_path is None and distil_teacher_name_or_path is None:
135144
# Large teacher is default
136145
param_dict["distil_teacher_name_or_path"] = "aloxatel/bert-base-mnli"
137146

138147
import nn_pruning.examples.text_classification.glue_sparse_xp as glue_sparse_xp
148+
139149
experiment = glue_sparse_xp.GlueSparseXP(param_dict)
140150

151+
# This does not actually use hyper parameter search right now, but it's useful for naming the output directory for example
141152
experiment.hyperparameter_search()
142153

143-
@cli.command()
144-
@click.pass_context
145-
@click.argument("path", default=None, type=click.Path(exists=True, resolve_path=True))
146-
@click.argument("output", default=None, type=click.Path(resolve_path=True))
147-
def command2(ctx, path, output):
148-
print(path + ":" + run2(path, output))
149-
150154

151155
def main():
152156
return cli()
153-
154-

0 commit comments

Comments
 (0)