Update app.py

app.py

@@ -1,13 +1,36 @@
 import streamlit as st
 import torch
-import transformers
+from transformers import AutoTokenizer, AutoModel, pipeline
+from torch import nn
 
 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():
-    return torch.load('model.pt'), AutoTokenizer.from_pretrained('bert-base-uncased')()
+    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()