Update app.py

app.py

@@ -7,11 +7,12 @@ from sklearn.preprocessing import LabelEncoder
 from keras.utils import pad_sequences
 from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
 
-st.markdown("### Hello, world!")
+st.markdown("### Paper category classification")
 st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
 # ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter
 
-text = st.text_area("TEXT HERE")
+title = st.text_area("TITLE HERE")
+abstract = st.text_area("ABSTRACT HERE")
 # ^-- показать текстовое поле. В поле text лежит строка, которая находится там в данный момент
 
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -20,63 +21,84 @@ model = BertForSequenceClassification.from_pretrained(
     num_labels = 44,)
 model.load_state_dict(torch.load("model_last_version.pt", map_location=torch.device('cpu')))
 MAX_LEN = 64
-tokens = tokenizer.encode_plus(text, add_special_tokens=True, max_length=MAX_LEN, truncation=True, padding='max_length')
-input_ids = torch.tensor(tokens['input_ids']).unsqueeze(0)
-attention_mask = torch.tensor(tokens['attention_mask']).unsqueeze(0)
+# Преобразуем название статьи в токены
+tokens = tokenizer(title, padding=True, truncation=True, return_tensors="pt")
 
+# Получаем предсказание модели для названия статьи и абстракта (если есть)
+input_ids = tokens['input_ids']
+attention_mask = tokens['attention_mask']
 logits = model(input_ids, attention_mask)[0]
-probs = torch.softmax(logits, dim=1)
-
-predicted_category = torch.argmax(probs).item()
 
 tags_names = ['acc-phys',
  'adap-org',
- "adap-org'",
+ "adap-org",
  'alg-geom',
  'astro-ph',
- "astro-ph'",
+ "astro-ph",
  'chao-dyn',
  'chem-ph',
  'cmp-lg',
- "cmp-lg'",
+ "cmp-lg",
  'comp-gas',
  'cond-mat',
- "cond-mat'",
+ "cond-mat",
  'cs',
  'dg-ga',
  'econ',
  'eess',
  'funct-an',
  'gr-qc',
- "gr-qc'",
+ "gr-qc",
  'hep-ex',
- "hep-ex'",
+ "hep-ex",
  'hep-lat',
- "hep-lat'",
+ "hep-lat",
  'hep-ph',
- "hep-ph'",
+ "hep-ph",
  'hep-th',
- "hep-th'",
+ "hep-th",
  'math',
  'math-ph',
  'mtrl-th',
  'nlin',
  'nucl-ex',
  'nucl-th',
- "nucl-th'",
+ "nucl-th",
  'patt-sol',
  'physics',
  'q-alg',
  'q-bio',
  'q-fin',
  'quant-ph',
- "quant-ph'",
+ "quant-ph",
  'solv-int',
  'stat']
+
+if abstract:
+    abstract_tokens = tokenizer(abstract, padding=True, truncation=True, return_tensors="pt")
+    abstract_input_ids = abstract_tokens['input_ids']
+    abstract_attention_mask = abstract_tokens['attention_mask']
+    abstract_logits = model(abstract_input_ids, abstract_attention_mask)[0]
+    logits += abstract_logits
+
+# Получаем вероятности и сортируем их в порядке убывания
+probs = torch.softmax(logits, dim=-1).squeeze()
+sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+
+# Считаем сумму вероятностей
+sum_probs = 0.0
+top_classes = []
+for i in range(len(sorted_probs)):
+    sum_probs += sorted_probs[i]
+    if sum_probs > 0.95 or sorted_probs[i] < 0.001:
+        break
+    top_classes.append((tag_names[sorted_indices[i].item()], sorted_probs[i].item()))
+
+# Выводим список тем с их вероятностями
 # from transformers import pipeline
 # pipe = pipeline("ner", "Davlan/distilbert-base-multilingual-cased-ner-hrl")
-raw_predictions = tags_names[predicted_category]#le.inverse_transform(prediction)#pipe(text)
+raw_predictions = top_classes#le.inverse_transform(prediction)#pipe(text)
 # тут уже знакомый вам код с huggingface.transformers -- его можно заменить на что угодно от fairseq до catboost
 
-st.markdown(f"{raw_predictions}")
+st.markdown(f"Possible categories for this article: {raw_predictions}")
 # выводим результаты модели в текстовое поле, на потеху пользователю