bert分类器

分类器用于将输入的字符串判断是传入训练的llama模型还是数据库。是一个二元分类且 1是数据库，0是llama（1是化工知识，0是日常问答）

在推理时和训练时传入的x有区别，因为训练过程中Dataloader会按照batch抽取数据且将他们分词结果处理成tensor形式。但是我们正常回答过程中，没有batch，也没有处理成tensor，所以需要改成如下形式（forward改就行）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch.nn as nn
from transformers import AutoModel


class MyNetwork(nn.Module):
    def __init__(self, args, freeze=False):
        super(MyNetwork, self).__init__()
        self.num_classes = args.num_classes#2
        self.bert_dim = args.bert_dim#768

        self.bert = AutoModel.from_pretrained(args.model_type)
        self.fc = nn.Linear(self.bert_dim, self.num_classes)

        if freeze:
            for name, param in self.bert.named_parameters():
                param.requires_grad = False

        # self.init_weights()

    def forward(self, x):
        # print(x)
        outputs = self.bert(torch.tensor(x['input_ids']).unsqueeze(0).cuda(), torch.tensor(x['attention_mask']).unsqueeze(0).cuda(), torch.tensor(x['token_type_ids']).unsqueeze(0).cuda())
        out = self.fc(outputs.last_hidden_state[:, 0, :])
        return out
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class t:
    def __init__(self):
        self.bert_dim = 768
        self.num_classes = 2
        self.model_type = '/data/lijiaqi/MultimodalAttack/BERT-Attack/model/bert-base-uncased'
        
args = t()
model = MyNetwork(args=args)
model.to('cuda')
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(args.model_type)
import torch
t = tokenizer('洛阳在哪个国家?', truncation=True, padding="max_length", max_length=128)
model(t)
#tensor([[0.8523, 0.2502]], device='cuda:0', grad_fn=<AddmmBackward0>)

VLLM 使用

1
2
3
4
5
from vllm import LLM

llm = LLM(model="/data/lijiaqi/chem/data/plm/output/20240315_chem_generalandchem_60000_merge",gpu_memory_utilization=0.4,swap_space=20)  # swap_space这个参数小的话容易爆。gpu_memory_utilization这个参数大的话容易爆。
output = llm.generate(["Hello, my name is"])
print(output[0].outputs[0].text)

整体

  • 加载分类器
  • 将网页的输入传入分类器
1
2
t = tokenizer('给定一张山脉的照片，创作一首有关景色的俳句。', truncation=True, padding="max_length", max_length=128,return_tensors="pt")#
model(t)
  • 根据分类