Update app.py

app.py

@@ -1,57 +1,53 @@
-import streamlit as st
-import pandas as pd
-import numpy as np
 import torch
+import torch.nn.functional as F
+import streamlit as st
 from transformers import BertTokenizer, BertForSequenceClassification
+import joblib
+
+# Загрузка модели, токенизатора и label_encoder
+model = torch.load("bert_model.pkl", map_location=torch.device('cpu'))
+tokenizer = joblib.load("bert_tokenizer.pkl")
+label_encoder = joblib.load("label_encoder.pkl")
+
+def predict_class_with_probabilities(text, model, tokenizer, label_encoder, max_len=128):
+    model.eval()
+
+    encodings = tokenizer(
+        text,
+        truncation=True,
+        padding="max_length",
+        max_length=max_len,
+        return_tensors="pt"
+    )
 
-# Загрузка модели и токенизатора с Hugging Face Hub
-model = torch.load('bert_model.pkl', map_location=torch.device('cpu'))
-tokenizer = joblib.load('bert_tokenizer.pkl')
-
-# Устройство для использования модели
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
-
-# Загрузка данных для поиска сходства
-try:
-    data = pd.read_excel('DATA_new.xlsx')
-    data_texts = data['Tags'].tolist()
-except FileNotFoundError:
-    st.error("Файл 'DATA_new.xlsx' не найден.")
-except Exception as e:
-    st.error(f"Ошибка загрузки файла: {e}")
-
-# Функция для нахождения сходства
-def find_similar_texts(input_text, top_n=5):
-    inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True).to(device)
     with torch.no_grad():
-        input_vector = model(**inputs).logits
-    data_vectors = []
-    for text in data_texts:
-        inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True).to(device)
-        with torch.no_grad():
-            data_vectors.append(model(**inputs).logits)
-    data_vectors = torch.stack(data_vectors).squeeze()
-    similarities = torch.nn.functional.cosine_similarity(input_vector, data_vectors)
-    similar_indices = torch.argsort(similarities, descending=True)[:top_n]
-    similar_texts = [data_texts[i] for i in similar_indices]
-    return similar_texts
-
-# Streamlit интерфейс в файле app.py
-if __name__ == "__main__":
-    st.title("Поиск сходства текстов")
-    st.write("Введите текст для поиска сходства")
-
-    input_text = st.text_area("Текст для поиска сходства")
-
-    if st.button("Найти похожие тексты"):
-        if input_text.strip():
-            similar_texts = find_similar_texts(input_text)
-            if similar_texts:
-                st.write("Похожие тексты:")
-                for text in similar_texts:
-                    st.write(f"- {text}")
-            else:
-                st.write("Нет похожих текстов для данного ввода.")
+        outputs = model(**encodings)
+        logits = outputs.logits
+
+        probabilities = F.softmax(logits, dim=1).squeeze().cpu().numpy()
+        predicted_class = torch.argmax(logits, dim=1).item()
+        predicted_label = label_encoder.inverse_transform([predicted_class])[0]
+
+    return predicted_label, probabilities
+
+def main():
+    st.title("Text Classification App with Hugging Face Space")
+    st.write("Введите текст, чтобы получить предсказание и вероятности классов.")
+
+    input_text = st.text_input("Введите текст для классификации:")
+
+    if st.button("Предсказать"):
+        if input_text:
+            predicted_class, probabilities = predict_class_with_probabilities(input_text, model, tokenizer, label_encoder)
+
+            st.write(f"**Предсказанный класс:** {predicted_class}")
+            st.write("**Вероятности для каждого класса:**")
+
+            for idx, prob in enumerate(probabilities):
+                class_label = label_encoder.inverse_transform([idx])[0]
+                st.write(f"{class_label}: {prob:.4f}")
         else:
-            st.error("Пожалуйста, введите текст для поиска сходства.")
+            st.write("Пожалуйста, введите текст для предсказания.")
+
+if __name__ == "__main__":
+    main()