app.py

import streamlit as st
import pandas as pd
import streamlit as st
import pickle
import time
from typing import Tuple
from sklearn.feature_extraction.text import TfidfVectorizer
import transformers
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score
import torch
from transformers import AutoTokenizer, AutoModel
from torch.utils.data import TensorDataset, DataLoader
from sklearn.preprocessing import LabelEncoder
import re
import string
import numpy as np
import torch.nn as nn
import json
import gensim
import torch.nn.functional as F
from transformers import GPT2LMHeadModel, GPT2Tokenizer


st.title('10-я неделя DS. Классификация отзывов, определение токсичности и генерация текста')

st.sidebar.header('Выберите страницу')
page = st.sidebar.radio("Выберите страницу", ["Вводная информация", "Классификация отзывов", "Классификация отзывов: неважная информация", "Определение токсичности", "Генерация текста"])

if page == "Вводная информация":
        
        st.subheader('*Задача №1*: Классификация отзывов на медицинские учреждения')
        st.write('Задача в двух словах: необходимо дать классификацию отзыва тремя моделями, время, за которое происходит классификаци отзыва, а также таблицу сравнения моделей по F-1 macro для моделей')

        st.subheader('*Задача №2*: Определение токсичности')
        st.write('Задача в двух словах: Оценка степени токсичности пользовательского сообщения ')

        st.subheader('*Задача №3*: Генерация текста')
        st.write('Задача в двух словах: Генерация текста GPT-моделью по пользовательскому prompt')

        st.subheader('☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️☀️')

        st.subheader('Выполнила команда "BERT": Алексей А., Светлана, Алиса')


if page == "Классификация отзывов":
    # Загрузка tf-idf модели и векторайзера
    with open('tf-idf/tf-idf.pkl', 'rb') as f:
        model_tf = pickle.load(f)

    with open('tf-idf/tf-idf_vectorizer.pkl', 'rb') as f:
        vectorizer_tf = pickle.load(f)

    # Загрузка словаря vocab_to_int и Word2Vec модели
    with open('lstm/vocab_to_int.json', 'r') as f:
        vocab_to_int = json.load(f)

    word2vec_model = gensim.models.Word2Vec.load("lstm/word2vec.model")

    stop_words = ['и', 'в', 'во', 'не', 'что', 'он', 'на', 'я', 'с', 'со', 'как', 'а', 'то', 'все', 'она', 'так', 'его', 'но', 'да', 'ты', 'к', 'у', 'же', 'вы', 'за', 'бы', 'по', 'только', 'ее', 'мне', 'было', 'вот', 'от', 'меня', 'еще', 'нет', 'о', 'из', 'ему', 'теперь', 'когда', 'даже', 'ну', 'вдруг', 'ли', 'если', 'уже', 'или', 'ни', 'быть', 'был', 'него', 'до', 'вас', 'нибудь', 'опять', 'уж', 'вам', 'ведь', 'там', 'потом', 'себя', 'ничего', 'ей', 'может', 'они', 'тут', 'где', 'есть', 'надо', 'ней', 'для', 'мы', 'тебя', 'их', 'чем', 'была', 'сам', 'чтоб', 'без', 'будто', 'чего', 'раз', 'тоже', 'себе', 'под', 'будет', 'ж', 'тогда', 'кто', 'этот', 'того', 'потому', 'этого', 'какой', 'совсем', 'ним', 'здесь', 'этом', 'один', 'почти', 'мой', 'тем', 'чтобы', 'нее', 'сейчас', 'были', 'куда', 'зачем', 'всех', 'никогда', 'можно', 'при', 'наконец', 'два', 'об', 'другой', 'хоть', 'после', 'над', 'больше', 'тот', 'через', 'эти', 'нас', 'про', 'всего', 'них', 'какая', 'много', 'разве', 'три', 'эту', 'моя', 'впрочем', 'хорошо', 'свою', 'этой', 'перед', 'иногда', 'лучше', 'чуть', 'том', 'нельзя', 'такой', 'им', 'более', 'всегда', 'конечно', 'всю', 'между']
    
    def data_preprocessing(text: str) -> str:
        text = text.lower()
        text = re.sub('<.*?>', '', text) # html tags
        text = ''.join([c for c in text if c not in string.punctuation])# Remove punctuation
        text = ' '.join([word for word in text.split() if word not in stop_words])
        text = [word for word in text.split() if not word.isdigit()]
        text = ' '.join(text)
        return text

    # Функция для предсказания класса отзыва
    def classify_review_tf(review):
        # Векторизация отзыва
        review_vector = vectorizer_tf.transform([review])
        # Предсказание
        start_time = time.time()
        prediction = model_tf.predict(review_vector)
        end_time = time.time()
        # Время предсказания
        prediction_time = end_time - start_time
        return prediction[0], prediction_time
    
    VOCAB_SIZE = len(vocab_to_int) + 1  # add 1 for the padding token
    EMBEDDING_DIM = 32
    HIDDEN_SIZE = 32
    SEQ_LEN = 100

    class BahdanauAttention(nn.Module):
        def __init__(self, hidden_size: torch.Tensor = HIDDEN_SIZE) -> None:
            super().__init__()

            self.W_q = nn.Linear(hidden_size, hidden_size)
            self.W_k = nn.Linear(hidden_size, hidden_size)
            self.V = nn.Linear(HIDDEN_SIZE, 1)

        def forward(
            self,
            keys: torch.Tensor,
            query: torch.Tensor
        ) -> Tuple[torch.Tensor, torch.Tensor]:
            query = self.W_q(query)
            keys = self.W_k(keys)

            energy = self.V(torch.tanh(query.unsqueeze(1) + keys)).squeeze(-1)
            weights = F.softmax(energy, -1)
            context = torch.bmm(weights.unsqueeze(1), keys)
            return context, weights
    
    embedding_matrix = np.zeros((VOCAB_SIZE, EMBEDDING_DIM))
    embedding_layer = torch.nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))

    class LSTMConcatAttention(nn.Module):
        def __init__(self) -> None:
            super().__init__()

            # self.embedding = nn.Embedding(VOCAB_SIZE, EMBEDDING_DIM)
            self.embedding = embedding_layer
            self.lstm = nn.LSTM(EMBEDDING_DIM, HIDDEN_SIZE, batch_first=True)
            self.attn = BahdanauAttention(HIDDEN_SIZE)
            self.clf = nn.Sequential(
                nn.Linear(HIDDEN_SIZE, 128),
                nn.Dropout(),
                nn.Tanh(),
                nn.Linear(128, 1)
            )
        
        def forward(self, x):
            embeddings = self.embedding(x)
            outputs, (h_n, _) = self.lstm(embeddings)
            att_hidden, att_weights = self.attn(outputs, h_n.squeeze(0))
            out = self.clf(att_hidden)
            return out, att_weights
        
    model_lstm = LSTMConcatAttention()  # Инициализируйте с теми же параметрами, что использовались при обучении
    model_lstm.load_state_dict(torch.load("lstm/lstm_model.pth"))
    model_lstm.eval()

        # Проверка и добавление токена <UNK>, если он отсутствует
    if '<UNK>' not in vocab_to_int:
        vocab_to_int['<UNK>'] = len(vocab_to_int)  # Присвоение нового уникального индекса

    # Проверка и добавление токена <PAD>, если он отсутствует
    if '<PAD>' not in vocab_to_int:
        vocab_to_int['<PAD>'] = len(vocab_to_int)  # Присвоение нового уникального индекса

    def text_to_vector(text, unknown_token_id=0):
            words = text.split()
            vector = [vocab_to_int.get(word, unknown_token_id) for word in words]  # здесь unknown_token_id - это ID для "неизвестных" слов
            return np.array(vector, dtype=np.int64)  # Убедитесь, что тип данных int64

    def classify_review_lstm(review):
        # Векторизация отзыва
        review_vector = text_to_vector(review)
        # Преобразование в тензор PyTorch и добавление размерности пакета (batch)
        review_tensor = torch.tensor(review_vector).unsqueeze(0)
        
        # Предсказание
        start_time = time.time()
        with torch.no_grad():
            prediction, _ = model_lstm(review_tensor)
        end_time = time.time()

        # Время предсказания
        prediction_time = end_time - start_time
        return prediction, prediction_time
    
    tokenizer_rubert = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
    model_rubert = AutoModel.from_pretrained("cointegrated/rubert-tiny2")
    clf_rubert = LogisticRegression(max_iter=1000)  # Предполагается, что ваша модель уже обучена

    with open('rubert/logistic_regression_model.pkl', 'rb') as f:
        clf_rubert = pickle.load(f)

    # Функция для предсказания
    def make_prediction(text):
        start_time = time.time()
        encoded = tokenizer_rubert(text, add_special_tokens=True, max_length=128, padding='max_length', truncation=True, return_tensors="pt")
        with torch.no_grad():
            outputs = model_rubert(**encoded)
        features = outputs.last_hidden_state[:, 0, :].numpy()
        prediction = clf_rubert.predict(features)
        end_time = time.time()
        prediction_time = end_time - start_time
        return prediction[0], prediction_time

    # Создание интерфейса Streamlit
    st.title('Классификатор отзывов на клиники')

    # Текстовое поле для ввода отзыва
    user_review = st.text_input('Введите ваш отзыв на клинику')

    if st.button('Классифицировать'):
        if user_review:
            # Классификация отзыва
            prediction_tf, pred_time_tf = classify_review_tf(user_review)
            st.write(f'Предсказанный класс TF-IDF: {prediction_tf}')
            st.write(f'Время предсказания TF-IDF: {pred_time_tf:.4f} секунд')
            prediction_lstm, pred_time_lstm = classify_review_lstm(user_review)
            st.write(f'Предсказанный класс LSTM: {prediction_tf}')
            st.write(f'Время предсказания LSTM: {pred_time_tf:.4f} секунд')
            prediction_rubert, pred_time_rubert = make_prediction(user_review)
            prediction_ru = 'negative' if prediction_rubert == 0 else 'positive'
            st.write(f'Предсказанный класс RuBERT: {prediction_ru}')
            st.write(f'Время предсказания RuBERT: {pred_time_rubert:.4f} секунд')
        else:
            st.write('Пожалуйста, введите отзыв')

if page == "Классификация отзывов: неважная информация":
# Создание данных для таблицы
    data = {
        "Название модели": ["TF-IDF", "LSTM", "RuBert tiny-2"],
        "F-1 macro score": ["0,94", "0,89", "0,90"]
    }

    # Создание DataFrame
    df = pd.DataFrame(data)

    # Отображение таблицы в Streamlit
    st.table(df)


if page == "Определение токсичности":

    # Функция для загрузки обученной модели
    def load_model(model_path):
        with open(model_path, 'rb') as file:
            model = pickle.load(file)
        return model

    # Загрузка обученной модели
    clf = load_model('toxic/logistic_regression_model_toxic.pkl')  # Укажите путь к файлу модели

    # Загрузка токенизатора и модели BERT
    tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny-toxicity")
    model = AutoModel.from_pretrained("cointegrated/rubert-tiny-toxicity")

    # Функция для предсказания токсичности сообщения
    def predict_toxicity(text):
        encoded = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
        with torch.no_grad():
            outputs = model(**encoded)
        features = outputs.last_hidden_state[:, 0, :].numpy()
        prediction = clf.predict(features)
        return prediction[0]

    # Создание интерфейса Streamlit
    st.title("Оценка токсичности сообщения")

    # Текстовое поле для ввода сообщения
    user_input = st.text_area("Введите сообщение для оценки")

    if st.button("Оценить"):
        if user_input:
            # Оценка токсичности сообщения
            prediction = predict_toxicity(user_input)
            if prediction > 0.5:
                st.write("Сообщение токсично")
                st.write(prediction)
            else:
                st.write("Сообщение не токсично")
                st.write(prediction)
        else:
            st.write("Пожалуйста, введите сообщение")

    
    
if page == "Генерация текста":
    # Путь к вашим весам модели
    model_weights_path = 'gpt-2/model.pt'

    # Загружаем токенизатор от GPT-2
    tokenizer = GPT2Tokenizer.from_pretrained("sberbank-ai/rugpt3small_based_on_gpt2")

    # Создаем экземпляр модели с архитектурой GPT-2, но без предварительно обученных весов
    model = GPT2LMHeadModel.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')

    # Загружаем веса вашей модели
    model.load_state_dict(torch.load(model_weights_path, map_location='cpu'))

    # Переносим модель на устройство (GPU или CPU)
    device = 'cpu'
    model.to(device)
    model.eval()

    def main():
        st.title("Генератор плохих отзывов больниц от ruGPT3")

        # Ввод текста от пользователя
        user_prompt = st.text_area("Введите текст-промпт:", "Я была в этой клинике..")

        # Виджеты для динамической регуляции параметров
        max_length = st.slider("Выберите max_length:", 10, 300, 100)
        temperature = st.slider("Выберите temperature:", 1.0, 10.0, step=0.2)
        top_k = st.slider("Выберите top_k:", 100, 500, 50)
        top_p = st.slider("Выберите top_p:", 0.1, 1.0, 0.95, step=0.05)
        num_beams = st.slider('Выберите num_beams:', 5, 40, step=1)

        # Генерация текста при нажатии на кнопку
        if st.button("Сгенерировать текст"):
            with torch.no_grad():
                prompt = tokenizer.encode(user_prompt, return_tensors='pt').to(device)
                out = model.generate(
                    input_ids=prompt,
                    max_length=max_length,
                    num_beams=num_beams,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    no_repeat_ngram_size=2,
                ).cpu().numpy()
                generated_text = tokenizer.decode(out[0], skip_special_tokens=True)
                st.subheader("Сгенерированный текст:")
                st.write(generated_text)
if __name__ == "__main__":
    main()