Update app.py

app.py

@@ -1,5 +1,36 @@
 import gradio as gr
 import pickle
+import numpy as np
+from torch import nn
+from transformers import AutoModelForSequenceClassification
+from sklearn.pipeline import Pipeline
+from skorch import NeuralNetClassifier
+from skorch.callbacks import LRScheduler, ProgressBar
+from skorch.hf import HuggingfacePretrainedTokenizer
+from torch.optim.lr_scheduler import LambdaLR
+from skorch.callbacks import EarlyStopping
+from sklearn.metrics import precision_recall_fscore_support
+from sklearn.metrics import balanced_accuracy_score
+from modAL.models import ActiveLearner
+from modAL.uncertainty import uncertainty_sampling
+
+
+class BertModule(nn.Module):
+    """ BERT model according to Skorch convention """
+    def __init__(self, name, num_labels):
+        super().__init__()
+        self.name = name
+        self.num_labels = num_labels
+        self.reset_weights()
+
+    def reset_weights(self):
+        self.bert = AutoModelForSequenceClassification.from_pretrained(
+            self.name, num_labels=self.num_labels
+        )
+
+    def forward(self, **kwargs):
+        pred = self.bert(**kwargs)
+        return pred.logits
 
 
 with open('learner.bin', 'rb') as f: