-
Notifications
You must be signed in to change notification settings - Fork 0
/
pruning_model.py
124 lines (111 loc) · 4.51 KB
/
pruning_model.py
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
import os
import torch
import transformers
from pathlib import Path
import pruning_dataset as dataset
class PruningModel(torch.nn.Module):
def __init__(self, config):
super(PruningModel, self).__init__()
self.config = config
# select model architecture
if config["architecture"] == "BART":
self.model = transformers.BartForConditionalGeneration.from_pretrained(
"facebook/bart-base"
)
self.tokenizer = transformers.BartTokenizerFast.from_pretrained("facebook/bart-base")
elif config["architecture"] == "T5":
self.model = transformers.T5ForConditionalGeneration.from_pretrained("t5-base")
self.tokenizer = transformers.T5TokenizerFast.from_pretrained("t5-base")
else:
raise Exception(
"Unknown architecture for module specified in config: currently, only T5-base (=T5) and BART-base (=BART) are supported."
)
def set_eval_mode(self):
"""Set model to eval mode."""
self.model.eval()
def save(self):
"""Save model."""
model_path = self.config["model_path"]
# create dir if not exists
model_dir = os.path.dirname(model_path)
Path(model_dir).mkdir(parents=True, exist_ok=True)
torch.save(self.model.state_dict(), model_path)
def load(self):
"""Load model."""
if torch.cuda.is_available():
state_dict = torch.load(self.config["model_path"])
else:
state_dict = torch.load(self.config["model_path"], torch.device('cpu'))
self.model.load_state_dict(state_dict)
if torch.cuda.is_available():
self.model = self.model.cuda()
def train(self, train_path, dev_path):
"""Train model."""
# load datasets
train_dataset = dataset.PruningDataset(self.config, self.tokenizer, train_path)
dev_dataset = dataset.PruningDataset(self.config, self.tokenizer, dev_path)
# arguments for training
training_args = transformers.Seq2SeqTrainingArguments(
output_dir="results", # output directory
num_train_epochs=self.config["num_train_epochs"], # total number of training epochs
per_device_train_batch_size=self.config[
"per_device_train_batch_size"
], # batch size per device during training
per_device_eval_batch_size=self.config[
"per_device_eval_batch_size"
], # batch size for evaluation
warmup_steps=self.config[
"warmup_steps"
], # number of warmup steps for learning rate scheduler
weight_decay=self.config["weight_decay"], # strength of weight decay
logging_dir="logs", # directory for storing logs
logging_steps=1000,
evaluation_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end="True"
# predict_with_generate=True
)
# create the object for training
trainer = transformers.Seq2SeqTrainer(
model=self.model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=dev_dataset,
)
# training progress
trainer.train()
# store model
self.save()
def inference(self, input):
"""Run the model on the given input."""
# encode
input_encodings = self.tokenizer(
input,
padding=True,
truncation=True,
max_length=self.config["max_input_length"],
return_tensors="pt",
)
if torch.cuda.is_available():
input_encodings = input_encodings.to(torch.device("cuda"))
# generate
output = self.model.generate(
input_ids=input_encodings["input_ids"],
attention_mask=input_encodings["attention_mask"],
no_repeat_ngram_size=self.config["no_repeat_ngram_size"],
num_beams=self.config["num_beams"],
early_stopping=self.config["early_stopping"],
)
# decoding
mentions_text = self.tokenizer.batch_decode(
output,
skip_special_tokens=True,
clean_up_tokenization_spaces=True,
)[0]
# obtain question words
mentions = self._get_mentions(mentions_text)
return mentions
def _get_mentions(self, mentions_text):
mentions = mentions_text.split("|")
mentions = [m.strip() for m in mentions]
return mentions