Update app.py

app.py

@@ -7,32 +7,34 @@ st.markdown("### Articles classificator.")
 # st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
 
 @st.cache
-def LoadModel():
+def get_bert_and_tokenizer():
     model_name = 'bert-base-uncased'
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    bert = AutoModel.from_pretrained(model_name)
-
-    class devops_model(nn.Module):
-        def __init__(self):
-            super(devops_model, self).__init__()
-            self.bert = bert
-            self.fc = nn.Sequential(
-                nn.Linear(768, 768),
-                nn.ReLU(),
-                nn.Dropout(0.3),
-                nn.BatchNorm1d(768),            
-                nn.Linear(768, 5),
-                nn.LogSoftmax(dim=-1)
-            )
-            
-        def forward(self, train_batch):
-            emb = self.bert(**train_batch)['pooler_output']
-            return self.fc(emb)
-
-    return torch.load('model.pt'), tokenizer
-
-
-model, tokenizer = LoadModel()
+    return AutoModel.from_pretrained(model_name), AutoTokenizer.from_pretrained(model_name)
+
+bert, tokenizer = get_bert_and_tokenizer()
+
+class devops_model(nn.Module):
+    def __init__(self):
+        super(devops_model, self).__init__()
+        self.bert = bert
+        self.fc = nn.Sequential(
+            nn.Linear(768, 768),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.BatchNorm1d(768),            
+            nn.Linear(768, 5),
+            nn.LogSoftmax(dim=-1)
+        )
+        
+    def forward(self, train_batch):
+        emb = self.bert(**train_batch)['pooler_output']
+        return self.fc(emb)
+
+@st.cache
+def LoadModel():
+    return torch.load('model.pt')
+
+model = LoadModel()
 
 def process(title, summary):
     text = title + summary