Update app.py

app.py

@@ -1,24 +1,21 @@
-import joblib
 import streamlit as st
 import pandas as pd
 import numpy as np
 import torch
-from transformers import BertTokenizer
+from transformers import BertTokenizer, BertForSequenceClassification
 
-# Загрузка модели и токенизатора с обработкой ошибки CUDA
-try:
-    model = torch.load('bert_model.pkl', map_location=torch.device('cpu'))
-except RuntimeError as e:
-    st.error(f"Ошибка загрузки модели: {e}")
-    model = None
-
-# Загрузка токенизатора BERT
+# Загрузка модели и токенизатора с 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['Text'].tolist()
+    data_texts = data['Tags'].tolist()
 except FileNotFoundError:
     st.error("Файл 'DATA_new.xlsx' не найден.")
 except Exception as e:
@@ -26,16 +23,12 @@ except Exception as e:
 
 # Функция для нахождения сходства
 def find_similar_texts(input_text, top_n=5):
-    if model is None:
-        st.error("Модель не загружена, невозможно выполнить поиск.")
-        return []
-    
-    inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True)
+    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)
+        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()