Update app.py

app.py

@@ -32,6 +32,20 @@ else:
     device = torch.device("cpu")
     print('No GPU available, using the CPU instead.')
 
+class BERTDataset(Dataset):
+    def __init__(self, dataset, sent_idx, label_idx, bert_tokenizer, vocab, max_len,
+                 pad, pair):
+        transform = BERTSentenceTransform(bert_tokenizer, max_seq_length=max_len,vocab=vocab, pad=pad, pair=pair)
+        #transform = nlp.data.BERTSentenceTransform(
+        #    tokenizer, max_seq_length=max_len, pad=pad, pair=pair)
+        self.sentences = [transform([i[sent_idx]]) for i in dataset]
+        self.labels = [np.int32(i[label_idx]) for i in dataset]
+
+    def __getitem__(self, i):
+        return (self.sentences[i] + (self.labels[i], ))
+
+    def __len__(self):
+        return (len(self.labels))
 
 # Kobert_softmax
 class BERTClassifier(nn.Module):