Update app.py

app.py

@@ -36,29 +36,7 @@ def get_top95(y_predict, convert_target):
 # Creating the customized model, by adding a drop out and a dense layer on top of distil bert to get the final output for the model. 
 from transformers import DistilBertModel, DistilBertTokenizer
 
-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'))
 # model.load_state_dict(checkpoint['model'])
 # optimizer.load_state_dict(checkpoint['opt']) 
@@ -90,6 +68,7 @@ def get_predict(title, abstract):
         attention_mask=inputs['attention_mask'],
     )
     logits = outputs[0]
+    print(logits)
     y_predict = torch.nn.functional.softmax(logits).cpu().detach().numpy()
     file_path = "sample.json"