Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-# Инициализация состояния сессии
+# Initialize session state
 if 'show_all' not in st.session_state:
     st.session_state.show_all = False
 if 'results' not in st.session_state:
@@ -146,9 +146,7 @@ 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')
@@ -159,25 +157,35 @@ try:
     )
 except OSError as e:
     st.error(f"Ошибка загрузки модели: {e}. Убедитесь, что модель доступна или укажите другую.")
-    st.stop()
+    st.stop()  # Остановка выполнения приложения при ошибке
+
 
 def classify_text(title, description):
-    text = f"{title} {description}"
-    topic_classifier = pipeline("text-classification", 
-                              model=model, 
-                              tokenizer=tokenizer, 
-                              top_k=len(id_to_cat))
+    """
+    Классифицирует текст и возвращает результаты в отсортированном виде.
+    Args:
+        title (str): Заголовок текста.
+        description (str): Краткое описание текста.
+    Returns:
+        list: Отсортированный список результатов классификации.
+    """
+    text = f"{title} {description}"  # Объединяем заголовок и описание
+    topic_classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k = len(id_to_cat))
     try:
-        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)
+        results = topic_classifier(text) 
     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")
 
@@ -185,43 +193,34 @@ 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("Идет классификация..."):
+        with st.spinner("Идет классификация..."):  # Индикатор загрузки
             st.session_state.results = classify_text(title, description)
-            st.session_state.show_all = False
+            st.session_state.show_all = False  # Reset to show only top 95%
 
-# Отображение результатов
+# Display results if available
 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()
+        st.subheader("Полные результаты классификации:")
+        for label, score in st.session_state.results:
+            st.write(f"- **{label}**: {score:.4f}")
     else:
-        if st.button("Показать все категории"):
+        st.subheader("Результаты классификации (top 95%):")
+        cumulative_prob = 0
+        for label, score in st.session_state.results:
+            st.write(f"- **{label}**: {score:.4f}")
+            cumulative_prob += score
+            if cumulative_prob >= 0.95:
+                break
+        
+        # Renamed button that refreshes the page
+        if st.button("Покажи все"):
             st.session_state.show_all = True
-            st.experimental_rerun()
+            st.experimental_rerun()  # Refresh the page
 
-# Отображение предупреждения только после первой попытки
-elif any([title, description]) and not st.session_state.results:
-    st.warning("Пожалуйста, нажмите кнопку 'Классифицировать'")
+elif title or description:  # небольшой костыль, чтобы при старте не было предупреждения
+    st.warning("Пожалуйста, нажмите кнопку 'Классифицировать'.")