Update app.py

app.py

@@ -67,28 +67,40 @@ infer_params = {"batch_size": INFER_BATCH_SIZE, "shuffle": False}
 infer_loader = DataLoader(infer_dataset, **infer_params)
 
 class BertClass(torch.nn.Module):
-  def __init__(self):
-    super(BertClass, self).__init__()
-    self.l1 = BertModel.from_pretrained(bert_path)
-    self.dropout = torch.nn.Dropout(HEAD_DROP_OUT)
-    self.classifier = torch.nn.Linear(768, 6)
-    
-  def forward(self, input_ids, attention_mask, token_type_ids):
-    output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
-    hidden_state = output_1[0]
-    pooler = hidden_state[:, 0]
-    pooler = self.dropout(pooler)
-    output = self.classifier(pooler)
-    return output
+    def __init__(self):
+        super(BertClass, self).__init__()
+        self.l1 = BertModel.from_pretrained(bert_path)
+        self.dropout = torch.nn.Dropout(HEAD_DROP_OUT)
+        self.classifier = torch.nn.Linear(768, 6)
+
+    # return_dict must equal False for Huggingface Transformers v4+
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        output_1 = self.l1(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            return_dict=False,
+        )
+        hidden_state = output_1[0]
+        pooler = hidden_state[:, 0]
+        pooler = self.dropout(pooler)
+        output = self.classifier(pooler)
+        return output
 
 class PretrainedBertClass(torch.nn.Module):
-  def __init__(self):
-    super(PretrainedBertClass, self).__init__()
-    self.l1 = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
-    
-  def forward(self, input_ids, attention_mask, token_type_ids):
-    output = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
-    return output
+    def __init__(self):
+        super(PretrainedBertClass, self).__init__()
+        self.l1 = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
+
+    # return_dict must equal False for Huggingface Transformers v4+
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        output = self.l1(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            return_dict=False,
+        )
+        return output
 
 # User selects model for front-end.
 option = st.selectbox("Select a text analysis model:", ("BERT", "Fine-tuned BERT"))
@@ -108,11 +120,12 @@ def inference():
             mask = data["mask"].to(device, dtype=torch.long)
             token_type_ids = data["token_type_ids"].to(device, dtype=torch.long)
             targets = data["targets"].to(device, dtype=torch.float)
-            outputs = model(ids, mask, token_type_ids, return_dict=False)
+            outputs = model(ids, mask, token_type_ids)
             final_targets.extend(targets.cpu().detach().numpy().tolist())
             final_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
     return final_outputs, final_targets
 
+
 prediction, targets = inference()
 prediction = np.array(prediction) >= 0.5
 targets = np.argmax(targets, axis=1)