Update app.py

app.py

@@ -1,53 +1,59 @@
-import torch
-import torch.nn.functional as F
 import joblib
-from transformers import BertTokenizer, BertForSequenceClassification
-
-# Загрузка предобученной модели и токенизатора
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+import streamlit as st
+import pandas as pd
+import numpy as np
+import torch
+from transformers import BertTokenizer
 
-model = torch.load('bert_model.pkl', map_location=device)
-model = model.to(device)
+# Загрузка модели и токенизатора с обработкой ошибки CUDA
+try:
+    model = torch.load('bert_model.pkl', map_location=torch.device('cpu'))
+except RuntimeError as e:
+    st.error(f"Ошибка загрузки модели: {e}")
 
+# Загрузка токенизатора BERT
 tokenizer = joblib.load('bert_tokenizer.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"
-    )
-
-    # Перенос токенизированных данных на устройство
-    encodings = {key: val.to(device) for key, val in encodings.items()}
-
+# Загрузка данных для поиска сходства
+try:
+    data = pd.read_excel('DATA_new.xlsx')
+    data_texts = data['Text'].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)
     with torch.no_grad():
-        # Предсказание модели
-        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
-
-# Пример использования функции
-input_text = "художественно разработка дизайн проекта"
-predicted_class, probabilities = predict_class_with_probabilities(input_text, model, tokenizer, label_encoder)
-
-# Вывод результатов
-print(f"Предсказанный класс: {predicted_class}")
-print("Вероятности для каждого класса:")
-for idx, prob in enumerate(probabilities):
-    class_label = label_encoder.inverse_transform([idx])[0]
-    print(f"{class_label}: {prob:.4f}")
+        input_vector = model(**inputs).logits
+    data_vectors = []
+    for text in data_texts:
+        inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
+        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("Нет похожих текстов для данного ввода.")
+        else:
+            st.error("Пожалуйста, введите текст для поиска сходства.")