Update app.py

app.py

@@ -37,7 +37,6 @@ def get_top95(y_predict, convert_target):
 from transformers import DistilBertModel, DistilBertTokenizer
 
 
-model = torch.load("pytorch_distilbert_news (3).bin", map_location=torch.device('cpu'))
 # model.load_state_dict(checkpoint['model'])
 # optimizer.load_state_dict(checkpoint['opt']) 
 # model.to("cpu")
@@ -50,6 +49,37 @@ model = torch.load("pytorch_distilbert_news (3).bin", map_location=torch.device(
 #     num_labels=8,
 #     return_dict=False)
 
+
+
+class DistillBERTClass(torch.nn.Module):
+    def __init__(self):
+        super(DistillBERTClass, self).__init__()
+        self.l1 = DistilBertModel.from_pretrained("distilbert-base-uncased")
+        self.pre_classifier = torch.nn.Linear(768, 768)
+        self.dropout = torch.nn.Dropout(0.3)
+        self.classifier = torch.nn.Linear(768, 8)
+
+    def forward(self, input_ids, attention_mask):
+        output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_state = output_1[0]
+        pooler = hidden_state[:, 0]
+        pooler = self.pre_classifier(pooler)
+        pooler = torch.nn.ReLU()(pooler)
+        pooler = self.dropout(pooler)
+        output = self.classifier(pooler)
+        return output
+
+
+model = DistillBERTClass()
+LEARNING_RATE = 1e-05
+
+optimizer = torch.optim.Adam(params =  model.parameters(), lr=LEARNING_RATE)
+
+
+
+model = torch.load("pytorch_distilbert_news (3).bin", map_location=torch.device('cpu'))
+
+
 def get_predict(title, abstract):
     tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased')
     # encoded_dict = tokenizer.encode_plus(