Update app.py

app.py

@@ -3,14 +3,33 @@ import torch
 
 @st.cache
 def Model():
-  from transformers import BertTokenizer, BertForSequenceClassification
-  model_name = "google/bert_uncased_L-4_H-256_A-4"
-  tokenizer = BertTokenizer.from_pretrained(model_name)
-  model = BertForSequenceClassification.from_pretrained(model_name, num_labels=8)
-  bn_state_dict = torch.load('model_w.pt')
+  from transformers import DebertaTokenizer, DebertaForSequenceClassification
+  tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-base")
+  model = DebertaForSequenceClassification.from_pretrained("microsoft/deberta-base", num_labels=8)
+  bn_state_dict = torch.load('model_weights.pt')
   model.load_state_dict(bn_state_dict)
-  return model
+  return model, tokenizer
 
+def Predict(model, tokenizer, text):
+  res = tokenizer(s, padding=True, truncation=True, return_tensors="pt", max_length=512)
+  #var.to("cuda:0")
+  res = model(**res)
+  logits = res.logits.softmax(dim=1)
+  logits = logits.numpy()[0]#logits.cpu().detach().numpy()[0]
+  return logits
+  
+def Print(logits, dictionary):
+  z = zip(logits, np.arange(0, 8))
+  z = sorted(z, key=lambda x: x[0], reverse=True)
+  sum, idx = 0, 0
+  while sum < 0.95:
+    st.markdown(f"{idx + 1}. ", dictionary[z[idx][1]])
+    sum += z[idx][0]
+    idx += 1
+
+def filter(title, abstract):
+  return True
+ 
 st.title('Классификация статьи по названию и описанию')
 # ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter
 
@@ -19,13 +38,13 @@ title = st.text_area("Введите название статьи:")
 abstract = st.text_area("Введите описание статьи:")
 # ^-- показать текстовое поле. В поле text лежит строка, которая находится там в данный момент
 
-from transformers import pipeline
-pipe = pipeline("ner", "Davlan/distilbert-base-multilingual-cased-ner-hrl")
-raw_predictions = pipe(title)
-# тут уже знакомый вам код с huggingface.transformers -- его можно заменить на что угодно от fairseq до catboost
+text = title + '. ' + abstract
+dictionary = ['computer science', 'economics', 'Electrical Engineering and Systems Science', 
+              'math', 'physics', 'quantitative biology', 'quantitative finance',
+              'statistics']
+if filter(title, abstract):
+  model, tokenizer = Model()
+  logits = Predict(model, tokenizer, text)
+  Print(logits, dictionary)
 
-st.markdown(f"{raw_predictions}")
 
-model = Model()
-st.markdown(f"{model}")
-# выводим результаты модели в текстовое поле, на потеху пользователю