基于NLP的恶意网页识别

基于NLP的恶意网页识别

引言

欢迎阅读《 基于NLP的恶意网页识别》，在前三篇中，我们已经使用PaddleNLP进行了恶意网页的分类，包括使用文本分类模型和预训练模型Fine-tune。本篇文章将着重优化模型，处理HTML标签提取结果不理想的情况，并最终将训练好的模型部署成可用的Python应用程序。

项目目录回顾

在前三篇文章中，我们已经完成了以下内容：

1. 使用PaddleNLP的文本分类模型进行简单的正常网页与被黑网页的二分类。
2. 使用PaddleNLP的预训练模型Fine-tune，提高HTML网页内容处理结果的判断准确率。
3. 进行正常网页与恶意网页的二分类，提取HTML标签信息判断网页是否正常。

本篇文章将以第四篇为基础，继续优化HTML标签提取结果，训练并评估模型，并最终将模型部署成可用的Python应用程序。

优化HTML标签提取结果

在之前的训练中，我们发现有些样本的HTML标签提取结果不够理想，主要集中在<script>标签内的信息清理不完整。为了解决这个问题，我们可以设计逻辑判断，将这部分内容留给下一个流程（比如提取文本信息或人工核验）。通过观察样本，我们发现这部分内容大多是恶意网页，所以即使略过也不会对模型的准确性产生太大影响。接下来，我们将优化标签提取结果。

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)

这段代码将原始的训练集和验证集中包含<script>标签内信息不完整的样本去除，得到新的训练集train_list2.txt和验证集eval_list2.txt。

使用预训练模型Fine-tune

接下来，我们使用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链接、图片信息等，加入到系统中，进一步提升恶意网页识别的准确性。