Models/bert_strim.py

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()

@st.cache_data
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]
    if prediction == 1:
        return "Позитивный отзыв 😀"
    else:
        return "Негативный отзыв 😟"

def bert_model_page():
    st.title("Классификация отзывов")
    user_input = st.text_area("Введите отзыв:")
    if st.button("Классифицировать"):
        if user_input:
            start_time = time.time()
            prediction = predict_sentiment(user_input)
            end_time = time.time()
            execution_time = end_time - start_time
            st.write("Результат классификации:", prediction)
            st.write(f'Время предсказания: {execution_time:.4f} секунд')
        else:
            st.write("Пожалуйста, введите отзыв для классификации.")