Update app.py

app.py

@@ -4,15 +4,13 @@ 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(allow_output_mutation=True)
-def get_bert_and_tokenizer():
-    model_name = 'bert-base-uncased'
-    # return AutoModel.from_pretrained(model_name), AutoTokenizer.from_pretrained(model_name)
+def get_tokenizer():
+    model_name = 'microsoft/deberta-v3-small'
     return AutoTokenizer.from_pretrained(model_name)
 
-tokenizer = get_bert_and_tokenizer()
+tokenizer = get_tokenizer()
 
 class devops_model(nn.Module):
     def __init__(self):
@@ -28,12 +26,12 @@ class devops_model(nn.Module):
         )
         
     def forward(self, train_batch):
-        emb = self.bert(**train_batch)['pooler_output']
+        emb = self.bert(**train_batch)['last_hidden_state'].mean(axis=1)
         return self.fc(emb)
 
 @st.cache
 def LoadModel():
-    return torch.load('model.pt', map_location=torch.device('cpu'))
+    return torch.load('model_full.pt', map_location=torch.device('cpu'))
 
 model = LoadModel()