欢迎阅读《 基于NLP的恶意网页识别》,在前三篇中,我们已经使用PaddleNLP进行了恶意网页的分类,包括使用文本分类模型和预训练模型Fine-tune。本篇文章将着重优化模型,处理HTML标签提取结果不理想的情况,并最终将训练好的模型部署成可用的Python应用程序。
在前三篇文章中,我们已经完成了以下内容:
本篇文章将以第四篇为基础,继续优化HTML标签提取结果,训练并评估模型,并最终将模型部署成可用的Python应用程序。
在之前的训练中,我们发现有些样本的HTML标签提取结果不够理想,主要集中在标签内的信息清理不完整。为了解决这个问题,我们可以设计逻辑判断,将这部分内容留给下一个流程(比如提取文本信息或人工核验)。通过观察样本,我们发现这部分内容大多是恶意网页,所以即使略过也不会对模型的准确性产生太大影响。接下来,我们将优化标签提取结果。
with open("train_list.txt", "r", encoding="utf-8") as f:
lines = f.readlines()
with open("train_list2.txt", "w", encoding="utf-8") as f_w:
for line in lines:
if ";" in line or "+" in line:
continue
f_w.write(line)
with open("eval_list.txt", "r", encoding="utf-8") as f:
lines = f.readlines()
with open("eval_list2.txt", "w", encoding="utf-8") as f_w:
for line in lines:
if ";" in line or "+" in line:
continue
f_w.write(line)
这段代码将原始的训练集和验证集中包含标签内信息不完整的样本去除,得到新的训练集
train_list2.txt
和验证集eval_list2.txt
。
接下来,我们使用PaddleNLP进行预训练模型Fine-tune,以优化HTML标签提取结果的预测准确率。首先,我们需要定义自定义数据集和加载预训练模型。
from paddlenlp.datasets import load_dataset
def read(data_path):
with open(data_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip('\n').split('\t')
words = ''.join(line[:-1])
labels = line[-1]
yield {'text': words, 'label': labels}
train_ds = load_dataset(read, data_path='train_list2.txt', lazy=False)
dev_ds = load_dataset(read, data_path='eval_list2.txt', lazy=False)
# 手动添加标签列表
train_ds.label_list = ['0', '1']
dev_ds.label_list = ['0', '1']
这段代码定义了自定义数据集,并手动添加了标签列表。接下来,我们加载预训练模型和tokenizer。
MODEL_NAME = "ernie-2.0-large-en"
ernie_model = paddlenlp.transformers.ErnieModel.from_pretrained(MODEL_NAME)
model = paddlenlp.transformers.ErnieForSequenceClassification.from_pretrained(MODEL_NAME, num_classes=len(train_ds.label_list))
tokenizer = paddlenlp.transformers.ErnieTokenizer.from_pretrained(MODEL_NAME)
我们使用PaddleNLP提供的数据处理和模型训练的接口,定义了数据处理函数和模型训练的参数。
from functools import partial
from paddlenlp.data import Stack, Tuple, Pad
from utils import convert_example, create_dataloader
batch_size = 128
max_seq_length = 64
trans_func = partial(
convert_example,
tokenizer=tokenizer,
max_seq_length=max_seq_length)
batchify_fn = lambda samples, fn=Tuple(
Pad(axis=0, pad_val=tokenizer.pad_token_id),
Pad(axis=0, pad_val=tokenizer.pad_token_type_id),
Stack(dtype="int64")): [data for data in fn(samples)]
train_data_loader = create_dataloader(
train_ds,
mode='train',
batch_size=batch_size,
batchify_fn=batchify_fn,
trans_fn=trans_func)
dev_data_loader = create_dataloader(
dev_ds,
mode='dev',
batch_size=batch_size,
batchify_fn=batchify_fn,
trans_fn=trans_func)
learning_rate = 5e-6
epochs = 5
warmup_proportion = 0.1
weight_decay = 0.1
num_training_steps = len(train_data_loader) * epochs
lr_scheduler = LinearDecayWithWarmup(learning_rate, num_training_steps, warmup_proportion)
optimizer = paddle.optimizer.AdamW(
learning_rate=lr_scheduler,
parameters=model.parameters(),
weight_decay=weight_decay,
apply_decay_param_fun=lambda x: x in [
p.name for n, p in model.named_parameters()
if not any(nd in n for nd in ["bias", "norm"])
])
criterion = paddle.nn.loss.CrossEntropyLoss()
metric = paddle.metric.Accuracy()
我们使用PaddleNLP提供的训练和评估接口,进行模型的训练与评估。同时,使用VisualDL进行可视化记录。
global_step = 0
for epoch in range(
1, epochs + 1):
with LogWriter(logdir="./visualdl") as writer:
for step, batch in enumerate(train_data_loader, start=1):
input_ids, segment_ids, labels = batch
logits = model(input_ids, segment_ids)
loss = criterion(logits, labels)
probs = F.softmax(logits, axis=1)
correct = metric.compute(probs, labels)
metric.update(correct)
acc = metric.accumulate()
global_step += 1
if global_step % 50 == 0:
print("global step %d, epoch: %d, batch: %d, loss: %.5f, acc: %.5f" % (
global_step, epoch, step, loss, acc))
writer.add_scalar(tag="loss", step=global_step, value=loss)
writer.add_scalar(tag="acc", step=global_step, value=acc)
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.clear_grad()
evaluate(model, criterion, metric, dev_data_loader)
model.save_pretrained('/home/aistudio/checkpoint')
tokenizer.save_pretrained('/home/aistudio/checkpoint')
训练完成后,我们可以将模型导出为静态图参数,以便后续部署使用。
state_dict = paddle.load('/home/aistudio/checkpoint/model_state.pdparams')
model.set_dict(state_dict)
model.eval()
model = paddle.jit.to_static(
model,
input_spec=[
paddle.static.InputSpec(
shape=[None, None], dtype="int64"),
paddle.static.InputSpec(
shape=[None, None], dtype="int64")
])
paddle.jit.save(model, '/home/aistudio/static_graph_params')
最后,我们将训练好的模型导出并进行部署,以便进行预测。这里我们假设已经准备好了一个HTML页面的内容,可以使用BeautifulSoup进行解析,提取HTML标签信息,然后使用训练好的模型进行预测。
html = BeautifulSoup(open('sample.html'), 'html.parser', from_encoding='utf-8')
def read_tags(text):
tags = []
class MyHTMLParser2(HTMLParser):
def handle_endtag(self, tag):
tags.append(tag)
parser = MyHTMLParser2()
parser.feed(text)
return tags
text = ','.join(read_tags(str(html.get_text)))
with open('sample.txt', 'w', encoding='utf-8') as f:
f.write(text)
!python predict.py --model_file=static_graph_params.pdmodel --params_file=static_graph_params.pdiparams
以上代码将HTML页面的标签信息提取并保存到sample.txt
文件中,然后使用训练好的模型进行预测。
在本文中,我们通过优化HTML标签提取结果,使用PaddleNLP进行预训练模型Fine-tune,最终将训练好的模型导出并部署成可用的Python应用程序。这一系列步骤构建了一个完整的恶意网页识别系统,可以帮助企业更好地保护用户免受网络攻击。在未来的工作中,我们可以考虑将网页内容的其他组成部分,如URL链接、图片信息等,加入到系统中,进一步提升恶意网页识别的准确性。