Update app.py

app.py

@@ -62,8 +62,7 @@ class MultiLabelDataset(set):
             'targets': torch.tensor(self.targets[idx], dtype=torch.float)
         }
 
-# Dataset and DataLoader
-trainSize = 0.4
+trainSize = 0.8
 trainData = new.sample(frac=trainSize,random_state=200)
 testData = new.drop(trainData.index).reset_index(drop=True)
 trainData = trainData.reset_index(drop=True)
@@ -74,17 +73,17 @@ testSet = MultiLabelDataset(testData, tokenizer, MAX_LEN)
 training_loader = DL(trainSet, batch_size=TRAIN_BATCH_SIZE, shuffle=True)
 testing_loader = DL(testSet, batch_size=VALID_BATCH_SIZE, shuffle=True)
 
-# model
-class DistilBERTClass(TNN.Module):
+# neural network
+class BERTClass(TNN.Module):
     def __init__(self):
-        super(DistilBERTClass, self).__init__()
+        super(BERTClass, self).__init__()
         self.l1 = BM.from_pretrained(modName)
         self.pre_classifier = TNN.Linear(768, 768)
         self.dropout = TNN.Dropout(0.1)
         self.classifier = TNN.Linear(768, 6)
 
     def forward(self, input_ids, attention_mask, token_type_ids):
-        out = self.l1(input_ids=input_ids, attention_mask=attention_mask)
+        out = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
         hidden_state = out[0]
         po = hidden_state[:, 0]
         po = self.pre_classifier(po)
@@ -93,7 +92,7 @@ class DistilBERTClass(TNN.Module):
         outs = self.classifier(po)
         return outs
 
-mod = DistilBERTClass()
+mod = BERTClass()
 mod.to(device)
 
 # Loss function and Optimizer