Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
-from transformers import BertTokenizer
+import torch.nn as nn
+from transformers import BertTokenizer, BertModel
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 names = ['负向', '正向']
@@ -8,6 +9,25 @@ names = ['负向', '正向']
 # 分词器
 tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
 
+class Model(nn.Module):
+    def __init__(self, bert_model):
+        super().__init__()
+        self.bert = bert_model
+        # 全连接，模型输入为768，分类为2
+        self.fc = nn.Linear(768, 2)
+
+    #
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        # 使用预训练模型提取特征, 上游任务不参与训练，锁定权重
+        with torch.no_grad():
+            # Correctly call the BertModel instance stored in self.bert
+            output = self.bert(input_ids, attention_mask, token_type_ids)
+        # 下游参与训练，二分类任务，获取最新后的状态
+        output = self.fc(output.last_hidden_state[:, 0])
+        # softmax激活函数，NV结构，获取特征值dim维度为1
+        output = output.softmax(dim=1)
+        return output
+
 # 加载预训练模型
 bert_model = BertModel.from_pretrained("ckiplab/bert-base-chinese").to(device)
 model = Model(bert_model).to(device)