app.py

import streamlit as st
import numpy as np
import pandas as pd
import time
import torch
import torch.nn as nn
from torch import tensor

import joblib
from dataclasses import dataclass
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, GPT2LMHeadModel, GPT2Tokenizer, GPT2Config
import json

from preprocessing import predict_review, data_preprocessing_hard
from model_lstm import LSTMClassifier
# from BERT_inputs import BertInputs

device = 'cpu'
classifier_bag = joblib.load('classifier_bag.pkl')
classifier_tf = joblib.load('classifier_tf.pkl')
BERT_lin_cl = joblib.load('BERT_base_model.pkl')

selected_model = st.sidebar.radio("Зачем пришел?", ("Классифиция отзывов лечебных учреждений", 
                                                    "Оценка степени токсичности пользовательского сообщения",
                                                    "Генерация текста GPT-моделью по пользовательскому prompt"))

# Классификация отзыва на поликлиники
if selected_model == "Классифиция отзывов лечебных учреждений":
    st.title("""
    Приложение классифицирует твой отзыв и подскажет, позитивный он или негативный
    """)
    st.write("""
    Классификация происходит с использованием классических ML моделей, нейросетевой модели LSTM, 
             и, как вариант, с использованием нейросетевой модели Bert-basic-ru для векторизации и линейной 
             регрессии для классификации.
    """)
    vectorizer_1 = joblib.load('vectorizer_bag.joblib')
    vectorizer_2 = joblib.load('vectorizer_tf.joblib')

    # LSTM
    with open('vocab_lstm.json', 'r') as file:
        vocab_to_int = json.load(file)

    @dataclass
    class ConfigRNN:
        vocab_size: int
        device : str
        n_layers : int
        embedding_dim : int
        hidden_size : int
        seq_len : int
        bidirectional : bool or int

    net_config = ConfigRNN(
        vocab_size = len(vocab_to_int)+1,
        device='cpu',
        n_layers=2,
        embedding_dim=64, 
        hidden_size=32,
        seq_len = 100,
        bidirectional=False)
    
    lstm = LSTMClassifier(net_config)
    lstm.load_state_dict(torch.load('lstm_model.pth', map_location=device))
    lstm.to(device)
    # lstm.eval()

    # BERT
    tokenizer = AutoTokenizer.from_pretrained("Geotrend/bert-base-ru-cased")
    model = AutoModel.from_pretrained("Geotrend/bert-base-ru-cased")
    # model.eval()
    MAX_LEN = 200

    data = pd.DataFrame({
    'Модель': ["BagOfWords", "TF-IDF", "LSTM", "BERT-based-ru"],
    'f1_macro': [0.934, 0.939, 0.831, 0.845]
    })

    st.subheader("""
    Немного информации о точности используемых моделей после обучения:
    """)
    # st.write(data)
    st.table(data)
    user_text_input = st.text_area('Введите ваш отзыв здесь:', '')

    if st.button('Предсказать'):
        start_time = time.time()
        predictions = []
        prediction_time = []
        X = vectorizer_1.transform([data_preprocessing_hard(user_text_input)])
        predictions.append(classifier_bag.predict(X))
        end_time = time.time()
        prediction_time.append(end_time - start_time)

        start_time = time.time()
        X1 = vectorizer_2.transform([data_preprocessing_hard(user_text_input)])
        predictions.append(classifier_tf.predict(X1))
        end_time = time.time()
        prediction_time.append(end_time - start_time)

        start_time = time.time()
        predictions.append(predict_review(model=lstm, review_text=user_text_input, net_config=net_config, 
                    vocab_to_int=vocab_to_int))
        end_time = time.time()
        prediction_time.append(end_time - start_time)

        tokens = tokenizer.encode(user_text_input, add_special_tokens=True)
        padded_tokens = tokens + [0] * (MAX_LEN - len(tokens))
        input_tensor = tensor(padded_tokens).unsqueeze(0)
        with torch.no_grad():
            outputs = model(input_tensor)
        X2 = outputs.last_hidden_state[:,0,:].detach().cpu().numpy()
        predictions.append(BERT_lin_cl.predict(X2))
        end_time = time.time()
        prediction_time.append(end_time - start_time)

        res = []
        for i in predictions:
            if i >= 0.5:
                res.append('Позитивный комментарий')
            else:
                res.append('Негативный комментарий')

        result = pd.DataFrame({
        'Модель': ["BagOfWords", "TF-IDF", "LSTM", "BERT-based-ru"],
        'Результат': res,
        'Время предсказания' : prediction_time
        })
        st.table(result)

        if len(result[result['Результат'] == 'Негативный комментарий']) >= 2:
            gif_url = 'https://i.gifer.com/LdC3.gif'
            st.image(gif_url, caption="Негативный коментарий")
            
        else:
            gif_url = 'https://media2.giphy.com/media/v1.Y2lkPTc5MGI3NjExOTdnYjJ1eTE0bjRuMGptcjhpdTk2YTYzeXEzMzlidWFsamY2bW8wZyZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/LUg1GEjapflW7Vg6B9/giphy.gif'
            st.image(gif_url, caption="Позитивный коментарий")
        


# Оценка степени токсичности пользовательского сообщения
elif selected_model == "Оценка степени токсичности пользовательского сообщения":
    st.title("""
    Приложение классифицирует, токсичный комментарий или нет
    """)

    st.write("""
    Классификация происходит с использованием нейросетевой модели rubert-tiny-toxicity.
    """)

    # Toxicity
    model_t_checkpoint = 'cointegrated/rubert-tiny-toxicity'
    tokenizer_t = AutoTokenizer.from_pretrained(model_t_checkpoint)
    model_t = AutoModelForSequenceClassification.from_pretrained(model_t_checkpoint)

    def text2toxicity(text, aggregate=True):
        with torch.no_grad():
            inputs = tokenizer_t(text, return_tensors='pt', truncation=True, padding=True).to(model_t.device)
            proba = torch.sigmoid(model_t(**inputs).logits).cpu().numpy()
        if isinstance(text, str):
            proba = proba[0]
        if aggregate:
            return 1 - proba.T[0] * (1 - proba.T[-1])
        return proba
    
    user_text_input = st.text_area('Введите ваш отзыв здесь:')

    if st.button('Предсказать'):
        start_time = time.time()
        proba = text2toxicity(user_text_input, True)
        end_time = time.time()
        prediction_time = end_time - start_time

        model_message = f'Предсказание модели:'
        if proba >= 0.5:
            st.write(f'Степень токсичности комментария {round(proba, 2)*100}%')
            gif_url = "https://media1.giphy.com/media/cInbau65cwPWUeGTIZ/giphy.gif?cid=6c09b952seqdtvky8yn2uq6bt3kvo1vu5sdzpkdznjvmtxsh&ep=v1_internal_gif_by_id&rid=giphy.gif&ct=s"
            st.image(gif_url, caption="ТОКСИК")
        else:
            st.write(f'Степень токсичности комментария {round(proba, 2)*100}%')
            gif_url = 'https://i.gifer.com/origin/51/518fbbf9cf32763122f9466d3c686bb3_w200.gif'
            st.image(gif_url, caption="МИЛОТА")
        st.write(f'Время предсказания: {prediction_time:.4f} секунд')
    


# Генерация текста GPT
elif selected_model == "Генерация текста GPT-моделью по пользовательскому prompt":
    st.title("""
    Нейросетевой гороскоп Якубовский-Дьяченко
    """)

    st.write("""
    Для генерации текста используется предобученная сеть GPT2. Дообучение проходило на гороскопах. 
    Общая длина текста для обучения 37 001 887 слов.
    """)
    user_text_input = st.text_area('Введите текст для формиорования гороскопа:')

    # GPT2
    path = 'model_pretra/'
    model = GPT2LMHeadModel.from_pretrained(path).to(device)
    tokenizer = GPT2Tokenizer.from_pretrained(path)
    
    temperature = st.slider('Градус дичи', 0.1, 2.0, 1.1, step=0.1)
    max_gen_length = st.slider('Максимальная длина генерации', 10, 500, 100, step=10)
    num_generations = st.slider('Количество генераций', 1, 10, 2, step=1)
    
    if st.button('Сделать гороскоп'):
        with st.spinner('Генерация текста...'):
            start_time = time.time()
            input_ids = tokenizer.encode(user_text_input, return_tensors="pt").to(device)
            model.eval()
            with torch.no_grad():
                out = model.generate(
                    input_ids,
                    do_sample=True,
                    num_beams=num_generations,
                    temperature=temperature,
                    top_p=0.9,
                    max_length=max_gen_length,
                )
            generated_text = tokenizer.decode(out[0], skip_special_tokens=True)
            end_time = time.time()
            prediction_time = end_time - start_time
    
            st.write(f'{generated_text}')