Update app.py

app.py

@@ -49,8 +49,7 @@ class ToxicModel(torch.nn.Module):
         self.dropout3 = nn.Dropout(0.3)
         self.dropout4 = nn.Dropout(0.4)
         self.dropout5 = nn.Dropout(0.5)
-        self.output = nn.Linear(config.hidden_size, NUM_CLASSES)
-        
+        self.output = nn.Linear(config.hidden_size, NUM_CLASSES)  
         
     def forward(self, input_ids, attention_mask, token_type_ids):
         transformer_out = self.transformer(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
@@ -64,7 +63,7 @@ class ToxicModel(torch.nn.Module):
         logits = (logits1 + logits2 + logits3 + logits4 + logits5) / 5
         return logits
  
- def inference_fn(model, input_ids=None, attention_mask=None, token_type_ids=None):  
+def inference_fn(model, input_ids=None, attention_mask=None, token_type_ids=None):  
     model.eval()
     input_ids = input_ids.to(device)  
     attention_mask = attention_mask.to(device)  
@@ -75,8 +74,7 @@ class ToxicModel(torch.nn.Module):
     out = output.sigmoid().detach().cpu().numpy().flatten()
        
     return out
-    
-    
+   
 def predict(comment=None) -> dict:  
     text = str(comment)
     text = " ".join(text.split())