Spaces:
Running
Running
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModel | |
| import torch | |
| from Models.bert_file import BERTClassifier | |
| import numpy as np | |
| import time | |
| tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2") | |
| model = BERTClassifier() | |
| device = 'cpu' | |
| model.load_state_dict(torch.load('Weights/BERTmodel_weights2.pth',map_location=torch.device('cpu'))) | |
| model.eval() | |
| def predict_sentiment(text): | |
| MAX_LEN = 100 | |
| encoded_review = tokenizer.encode_plus( | |
| text, | |
| max_length=MAX_LEN, | |
| add_special_tokens=True, | |
| return_token_type_ids=False, | |
| pad_to_max_length=True, | |
| return_attention_mask=True, | |
| return_tensors='pt', | |
| ) | |
| input_ids = encoded_review['input_ids'].to(device) | |
| attention_mask = encoded_review['attention_mask'].to(device) | |
| with torch.no_grad(): | |
| output = model(input_ids, attention_mask) | |
| prediction = torch.round(output).cpu().numpy()[0][0] | |
| return prediction | |
| def bert_model_page(): | |
| st.title("Классификация отзывов") | |
| user_input = st.text_area("Введите отзыв:") | |
| if st.button("Классифицировать"): | |
| start_time = time.time() | |
| prediction = predict_sentiment(user_input) | |
| end_time = time.time() | |
| execution_time = end_time - start_time | |
| if prediction > 0.5: | |
| st.write("Отзыв положительный 🌟") | |
| else: | |
| st.write("Отзыв отрицательный 😞") | |
| st.write(f'Время предсказания: {execution_time:.4f} секунд') | |