def main():
sents_train, labels_train = loading_dataset(train_dataset_path)
sents_dev, labels_dev = loading_dataset(dev_dataset_path)
labels = ["0", "1"]
print("Labels: ", labels) # [0, 1]
# FastBERT model
model = FastBERT(kernel_name="uer_bert_tiny_zh",
labels=labels,
device="cuda:0" if torch.cuda.is_available() else "cpu")
# # FastGPT model
# model = FastGPT(
# kernel_name="uer_gpt_zh",
# labels=labels,
# device="cuda:0" if torch.cuda.is_available() else "cpu"
# )
model.fit(
sents_train,
labels_train,
sentences_dev=sents_dev,
labels_dev=labels_dev,
finetuning_epochs_num=3,
distilling_epochs_num=5,
report_steps=100,
model_saving_path=model_saving_path,
verbose=True,
)
def main():
sents_test, labels_test = loading_dataset(test_dataset_path)
samples_num = len(sents_test)
model = FastBERT(kernel_name="uer_bert_tiny_zh",
labels=labels,
device="cuda:0" if torch.cuda.is_available() else "cpu")
model.load_model(model_path)
correct_num = 0
exec_layer_list = []
for sent, label in zip(sents_test, labels_test):
label_pred, exec_layer = model(sent, speed=speed)
if label_pred == label:
correct_num += 1
exec_layer_list.append(exec_layer)
acc = correct_num / samples_num
ave_exec_layers = np.mean(exec_layer_list)
print("Acc = {:.3f}, Ave_exec_layers = {}".format(acc, ave_exec_layers))