Update app.py

app.py

@@ -1,6 +1,12 @@
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
+# Инициализация состояния сессии
+if 'show_all' not in st.session_state:
+    st.session_state.show_all = False
+if 'results' not in st.session_state:
+    st.session_state.results = []
+
 id_to_cat = {0: 'Performance',
  1: 'Molecular Networks',
  2: 'Operating Systems',
@@ -140,7 +146,9 @@ id_to_cat = {0: 'Performance',
  136: 'Nuclear Experiment',
  137: 'Artificial Intelligence'}
 
-# Загружаем модель (замените на вашу модель, если нужно)
+id_to_cat = {0: 'Performance', ...}  # Ваш полный словарь категорий
+
+# Загружаем модель
 model_name = 'checkpoint'
 try:
     tokenizer = AutoTokenizer.from_pretrained('distilbert-base-cased')
@@ -151,42 +159,25 @@ try:
     )
 except OSError as e:
     st.error(f"Ошибка загрузки модели: {e}. Убедитесь, что модель доступна или укажите другую.")
-    st.stop()  # Остановка выполнения приложения при ошибке
-
+    st.stop()
 
 def classify_text(title, description):
-    """
-    Классифицирует текст и возвращает результаты в отсортированном виде.
-
-    Args:
-        title (str): Заголовок текста.
-        description (str): Краткое описание текста.
-        show_all (bool): Показывать ли все результаты, независимо от порога.
-        threshold (float): Порог суммарной вероятности.
-
-    Returns:
-        list: Отсортированный список результатов классификации.
-    """
-    text = f"{title} {description}"  # Объединяем заголовок и описание
-    topic_classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k = len(id_to_cat))
+    text = f"{title} {description}"
+    topic_classifier = pipeline("text-classification", 
+                              model=model, 
+                              tokenizer=tokenizer, 
+                              top_k=len(id_to_cat))
     try:
-        results = topic_classifier(text) 
-        # results = topic_classifier(text, candidate_labels, multi_label=True)  # multi_label=True для нескольких меток
+        results = topic_classifier(text)
+        processed = []
+        for item in results[0]:
+            label_id = int(item['label'].split('_')[1])
+            processed.append((id_to_cat[label_id], item['score']))
+        return sorted(processed, key=lambda x: x[1], reverse=True)
     except Exception as e:
         st.error(f"Ошибка классификации: {e}")
         return []
 
-    for i in results[0]:
-        i['label'] = id_to_cat[int(i['label'].split('_')[1])]
-
-
-    filtered_results = []
-    for i in results[0]:
-        filtered_results.append((i['label'], i['score']))
-    return filtered_results
-        
-
-
 # --- Интерфейс Streamlit ---
 st.title("Классификация статей 1")
 
@@ -194,29 +185,43 @@ st.title("Классификация статей 1")
 title = st.text_input("Заголовок статьи")
 description = st.text_area("Краткое описание статьи", height=150)
 
-# Кнопка "Классифицировать"
+# Кнопка классификации
 if st.button("Классифицировать"):
     if not title and not description:
         st.warning("Пожалуйста, заполните хотя бы одно поле.")
     else:
-        with st.spinner("Идет классификация..."):  # Индикатор загрузки
-            results = classify_text(title, description)
-            if results:
-              st.subheader("Результаты классификации (top 95%):")
-              cumulative_prob = 0
-              for label, score in results:
-                  st.write(f"- **{label}**: {score:.4f}")
-                  cumulative_prob += score
-                  if cumulative_prob >= 0.95:
-                      break
+        with st.spinner("Идет классификация..."):
+            st.session_state.results = classify_text(title, description)
+            st.session_state.show_all = False
 
-              # Кнопка "Показать все"
-              if st.button("Показать все категории"):
-                  st.subheader("Полные результаты классификации:")
-                  for label, score in results:
-                      st.write(f"- **{label}**: {score:.4f}")
-            else:
-                st.info("Не удалось получить результаты классификации.")
+# Отображение результатов
+if st.session_state.results:
+    st.subheader("Результаты классификации:")
+    
+    # Определение порога отображения
+    cumulative = 0
+    shown_results = []
+    for label, score in st.session_state.results:
+        if not st.session_state.show_all and cumulative < 0.95:
+            shown_results.append((label, score))
+            cumulative += score
+        else:
+            shown_results.append((label, score))
+    
+    # Отображение результатов
+    for label, score in shown_results:
+        st.write(f"- **{label}**: {score:.4f}")
+    
+    # Кнопка переключения режима отображения
+    if st.session_state.show_all:
+        if st.button("Скрыть подробности"):
+            st.session_state.show_all = False
+            st.experimental_rerun()
+    else:
+        if st.button("Показать все категории"):
+            st.session_state.show_all = True
+            st.experimental_rerun()
 
-elif title or description: #небольшой костыль, чтобы при старте не было предупреждения
-    st.warning("Пожалуйста, заполните хотя бы одно поле.")
+# Отображение предупреждения только после первой попытки
+elif any([title, description]) and not st.session_state.results:
+    st.warning("Пожалуйста, нажмите кнопку 'Классифицировать'")