app.py

import streamlit as st 
import torch
from transformers import AutoTokenizer, AutoModel
import joblib
import numpy as np
import json
import pandas as pd
from gensim.models import Word2Vec
import time
from typing import Tuple
from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F
import torch.nn as nn
from torchmetrics import Accuracy
from torchmetrics.functional import f1_score
from string import punctuation
import sklearn as sk
import os
from sklearn.linear_model import LogisticRegression




def main():
    # Боковая панель с навигацией
    menu = ["Классификация отзыва на рестораны", "Классификация тематики новостей из телеграмм каналов", "Генерация текста GPT-моделью"]
    choice = st.sidebar.radio("Навигация", menu)

    # Отображение контента в зависимости от выбранной страницы
    if choice == "Классификация отзыва на рестораны":
        page_restoran()
    elif choice == "Классификация тематики новостей из телеграмм каналов":
        page_telegramm()
    elif choice == "Генерация текста GPT-моделью":
        page_generaition()

with open('vocab_lstm_att.json', 'r') as fp:
    vocab_to_int = json.load(fp)  

def page_restoran():
    model_path = "word2vec_model.bin"

    if os.path.exists(model_path):
        wv = Word2Vec.load(model_path)
    else:
        print(f"File '{model_path}' not found.")
        model_path = "word2vec_model.bin"
    wv = Word2Vec.load(model_path)

    VOCAB_SIZE = len(vocab_to_int)+1
    HIDDEN_SIZE = 128
    SEQ_LEN = 128
    DEVICE='cpu'
    EMBEDDING_DIM = 128 

    embedding_matrix = torch.load('embedding_matrix.pt')
    embedding_layer = torch.nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))

    class ConcatAttention(nn.Module):
        def __init__(
                self, 
                hidden_size: int = HIDDEN_SIZE
                ) -> None:
        
            super().__init__()
            self.hidden_size = hidden_size
            self.linear = nn.Linear(hidden_size, hidden_size)
            self.align  = nn.Linear(hidden_size * 2, hidden_size)
            self.tanh   = nn.Tanh()

        def forward(
                self, 
                lstm_outputs: torch.Tensor, 
                final_hidden: torch.Tensor 
                ) -> Tuple[torch.Tensor, torch.Tensor]:
        
            att_weights = self.linear(lstm_outputs)
            att_weights = torch.bmm(att_weights, final_hidden.unsqueeze(2))
            att_weights = F.softmax(att_weights.squeeze(2), dim=1)
            cntxt       = torch.bmm(lstm_outputs.transpose(1, 2), att_weights.unsqueeze(2))
            concatted   = torch.cat((cntxt, final_hidden.unsqueeze(2)), dim=1)
            att_hidden  = self.tanh(self.align(concatted.squeeze(-1)))
            return att_hidden, att_weights



    class LSTMConcatAttention(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.embedding = embedding_layer
            self.lstm = nn.LSTM(EMBEDDING_DIM, HIDDEN_SIZE, batch_first=True)
            self.attn = ConcatAttention(HIDDEN_SIZE)
            self.clf = nn.Sequential(
                nn.Linear(HIDDEN_SIZE, 128),
                nn.Dropout(),
                nn.Tanh(),
                nn.Linear(128, 6)
            )
    
        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_att = LSTMConcatAttention()
    model_lstm_att.load_state_dict(torch.load('lstm_att_model.pt', map_location='cpu'))
    model_lstm_att.eval()


    def pred(text): 
        start_time = time.time()
        text = text.lower()
        text = ''.join([c for c in text if c not in punctuation])
        text = [vocab_to_int[word] for word in text.split() if vocab_to_int.get(word)]
        if len(text) <= 128:
            zeros = list(np.zeros(128 - len(text)))
            text = zeros + text
        else:
            text = text[: 128]
        text = torch.Tensor(text)
        text = text.unsqueeze(0)
        text = text.type(torch.LongTensor)
        pred = model_lstm_att(text)[0].argmax(1)
        labels = {0: '0', 1:'1', 2:'2', 3: '3', 4:'4', 5:'5'}
        end_time = time.time()
        inference_time = end_time - start_time
    
        return f"***{labels[pred.item()]}***, время предсказания: ***{inference_time:.4f} сек***."


    """
    ## Классификация отзывов на ресторан
    """
    st.image('2024-02-02 12.58.31.jpg')
    st.title('Напиши отзыв на ресторан🌹')
    user_input = st.text_area('Введите отзыв', '')

    if st.button('Классифицировать LSTM!🤩'):
        if user_input.strip() == '':
            st.error('Введите, пожалуйста, отзыв!')
        else:
            st.success(pred(user_input))

    tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
    model = AutoModel.from_pretrained("cointegrated/rubert-tiny2")

    # Загрузка предобученной модели логистической регрессии
    lr_model = joblib.load('logistic_regression_model1.pkl') # веса модели

    # Расшифровка оценок
    class_mapping = {
        1: 'Не удовлетворительно',
        2: 'Удовлетворительно',
        3: 'Хорошо',
        4: 'Отлично',
        0: 'Превосходно'
    }

        # Функция для классификации текста
    def embed_destil(text, model, tokenizer):
        t = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
        with torch.no_grad():
            model_output = model(**{k: v.to(model.device) for k, v in t.items()})
        embeddings = model_output.last_hidden_state[:, 0, :]
        embeddings = torch.nn.functional.normalize(embeddings)
        return embeddings
    

    if st.button('Классифицировать ruBERT!🤩'):
        if user_input.strip() == '':
            st.error('Пожалуйста, напишите текст отзыва')
        else:
            X= embed_destil(user_input, model, tokenizer)
            predictions= lr_model.predict(X)
            res=class_mapping[predictions[0]]
            st.success(f'Оценка заведению: {res}') 
    model_ml = LogisticRegression()
    vectorizer = joblib.load("vectorizer.pkl")


    def preprocess(text):
    # Убедитесь, что text - это список
        if isinstance(text, str):
            text = [text]
    # Преобразуйте текст
        text = vectorizer.transform(text)
        return text

    model = model_ml
    model = joblib.load("log_reg_tfidf_model.pkl")

    def predict(text):
        start_time = time.time()
        text = preprocess(text)
        predicted_label = model.predict(text)
        dict = {0: '0', 1:'1', 2:'2', 3: '3', 4:'4', 5:'5'}
        predicted_label_text = dict[predicted_label[0]]
        end_time = time.time()

        inference_time = end_time - start_time

        return f"***{predicted_label_text}***, время предсказания: ***{inference_time:.4f} сек***."


    if st.button('Классифицировать ML-TFIDF!🤩'):
        if user_input.strip() == '':
            st.error('Введите, пожалуйста, отзыв!')
        else:
            st.success(predict(user_input))      
        
# if choice_model == 'LSTM(attention)':
#     if text:
#         st.write(pred(text))


    data = pd.DataFrame({'Модель': ['TFIDF-LogReg', 'LSTM(Attention)', 'BERT model'], 'F1-macro': [0.25, 0.2, 0.261]})
    # Вывод таблицы
    checkbox = st.sidebar.checkbox("Таблица f1-macro")
    if checkbox:
        st.write("<h1 style='text-align: center; font-size: 20pt;'>Оценка качества моделей по метрике f1-macro</h1>", unsafe_allow_html=True)
        st.table(data)
    


def page_telegramm():
    st.header("Определить тему новости 📑")
    # Добавьте контент для страницы 2
    st.image('collage.jpg')

    # Загрузка предобученной модели RuBERT
    tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
    model = AutoModel.from_pretrained("cointegrated/rubert-tiny2")

    # Загрузка предобученной модели логистической регрессии
    lr_model = joblib.load('trained_logistic_regression_model.pkl') 

    class_mapping = {
        1: 'Мода🛍️',
        2: 'Спорт🏀',
        3: 'Технологии📱',
        4: 'Финансы💰',
        0: 'Криптовалюта💸'
    }   

    # Функция для классификации текста
    def embed_bert_cls(text, model, tokenizer):
        t = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
        with torch.no_grad():
            model_output = model(**{k: v.to(model.device) for k, v in t.items()})
        embeddings = model_output.last_hidden_state[:, 0, :]
        embeddings = torch.nn.functional.normalize(embeddings)
        return embeddings

    user_input = st.text_area('Введите текст для классификации:', '')

    if st.button('Классифицировать!🤩'):
        if user_input.strip() == '':
            st.error('Пожалуйста, введите текст!😡')
        else:
            X= embed_bert_cls(user_input, model, tokenizer)
            predictions= lr_model.predict(X)
            res=class_mapping[predictions[0]]
            st.success(f'Я думаю, что тема новости: {res}')



def page_generaition():
    st.subheader("Генерация текста GPT-моделью")
    # Добавьте контент для страницы 3
       


if __name__ == "__main__":
    main()