Update app.py

app.py

@@ -6,11 +6,12 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 
 st.markdown('## Классификатор статей')
-st.write('Данный сервис предназначен для выбора темы статьи,     \n' \
-    'основываясь на ее названии и краткой выжимкой текста статьи.       \n' \
-    'Сервис работает благодаря fine-tune версии модели distil bert.     \n' \
+st.write('Данный сервис предназначен для выбора темы статьи [по таксономии arxiv.org](https://arxiv.org/category_taxonomy),     \n' \
+    'основываясь на ее названии и краткой выжимки текста статьи.       \n' \
+    'Сервис работает благодаря fine-tune версии модели [distil bert](https://huggingface.co/distilbert/distilbert-base-cased) [1].     \n' \
     'Данные для обучения были взяты [отсюда](https://www.kaggle.com/datasets/neelshah18/arxivdataset).       \n' \
-    'Поддерживается ввод только английского языка.')
+    'Поддерживается ввод только английского языка.      \n')
+
 st.markdown('#### Введите название статьи и ее краткое содержание:')
 
 device = torch.device('cpu')
@@ -60,7 +61,7 @@ def load_model():
     model_name = 'model'
     cat_count = 358
 
-    checkpoint = torch.load(os.path.join(chkp_folder, f"{model_name}.pt"), weights_only=False, map_location=torch.device('cpu'))
+    checkpoint = torch.load(os.path.join(chkp_folder, f"{model_name}.pt"), weights_only=False, map_location=device)
 
     # Создаём те же классы, что и внутри чекпоинта
 
@@ -92,6 +93,12 @@ case_['summary'] = st.text_area("Краткое содержание:", value=""
 
 if case_['title'] or case_['summary']:
     categories, probabilities = predict_category(case_, model, tokenizer)
-    st.write('Возможные категории:')
+    st.markdown('#### Возможные категории:')
     for i, cat in enumerate(categories):
-        st.write(f'{ind_to_cat[cat]}')
+        st.markdown("- " + f'{ind_to_cat[cat]}')
+
+st.write(
+    '''[1] DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter,
+    Victor Sanh and Lysandre Debut and Julien Chaumond and Thomas Wolf,
+    ArXiv, 2019, abs/1910.01108'''
+)