fgeegdfr

Владик/model.py

@@ -0,0 +1,229 @@
+from PIL import Image, ImageFilter, ImageDraw
+import streamlit as st
+
+import pickle
+import numpy as np
+import torch
+import torch.nn as nn
+from torch import Tensor
+from dataclasses import dataclass
+from typing import Union
+import re
+import string
+import pymorphy3
+import nltk
+from nltk.corpus import stopwords
+# stop_words = set(stopwords.words("english"))
+stop_words = nltk.download('stopwords')
+
+
+# ------------------------------------------------------------#
+# Упрощенный метод создания класса
+
+@dataclass
+class ConfigRNN:
+    vocab_size: int  # сколько слов - столько embedding-ов; для инициализации embedding параметров
+    device: str
+    n_layers: int
+    embedding_dim: int  # чем больше, тем сложнее можно закодировать слово
+    hidden_size: int
+    seq_len: int
+    bidirectional: Union[bool, int]
+
+
+net_config = ConfigRNN(
+    vocab_size=17259 + 1,  # -> hand
+    device="cpu",
+    n_layers=1,
+    embedding_dim=8,  # не лучшее значение, но в рамках задачи сойдет
+    hidden_size=16,
+    seq_len=30,  # -> hand
+    bidirectional=False,
+)
+# ------------------------------------------------------------#
+
+
+class LSTMClassifier(nn.Module):
+    def __init__(self, rnn_conf=net_config) -> None:
+        super().__init__()
+
+        self.embedding_dim = rnn_conf.embedding_dim
+        self.hidden_size = rnn_conf.hidden_size
+        self.bidirectional = rnn_conf.bidirectional
+        self.n_layers = rnn_conf.n_layers
+
+        self.embedding = nn.Embedding(rnn_conf.vocab_size, self.embedding_dim)
+        self.lstm = nn.LSTM(
+            input_size=self.embedding_dim,
+            hidden_size=self.hidden_size,
+            bidirectional=self.bidirectional,
+            batch_first=True,
+            num_layers=self.n_layers,
+            dropout=0.5
+        )
+        self.bidirect_factor = 2 if self.bidirectional else 1
+        self.clf = nn.Sequential(
+            nn.Linear(self.hidden_size * self.bidirect_factor, 32),
+            nn.Dropout(),
+            nn.Tanh(),
+            nn.Dropout(),
+            nn.Linear(32, 5)  # len(df['label'].unique())
+        )
+
+    def model_description(self):
+        direction = "bidirect" if self.bidirectional else "onedirect"
+        return f"lstm_{direction}_{self.n_layers}"
+
+    def forward(self, x: torch.Tensor):
+        embeddings = self.embedding(x)
+        out, _ = self.lstm(embeddings)
+        # print(out.shape)
+        # [все элементы батча, последний h_n, все элементы последнего h_n]
+        out = out[:, -1, :]
+        # print(out.shape)
+        out = self.clf(out)
+        return out
+# ------------------------------------------------------------#
+# Загрузка модели
+
+
+@st.cache_resource
+def load_model():
+    model = LSTMClassifier(net_config)
+    model.load_state_dict(torch.load(
+        "models/lstm_weights.pth", map_location=torch.device("cpu")))
+    model.eval()
+    return model
+
+
+model_lstm = load_model()
+# ------------------------------------------------------------#
+
+
+def padding(text_int: list, seq_len: int) -> np.ndarray:
+    """Make left-sided padding for input list of tokens
+
+    Args:
+        review_int (list): input list of tokens
+        seq_len (int): max length of sequence, it len(review_int[i]) > seq_len it will be trimmed, else it will be padded by zeros
+
+    Returns:
+        np.array: padded sequences
+    """
+    features = np.zeros((len(text_int), seq_len), dtype=int)
+    for i, review in enumerate(text_int):
+        if len(review) <= seq_len:
+            zeros = list(np.zeros(seq_len - len(review)))
+            new = zeros + review
+        else:
+            new = review[:seq_len]
+        features[i, :] = np.array(new)
+    return features
+
+
+morph = pymorphy3.MorphAnalyzer()
+
+
+def lemmatize(text):
+    # Разбиваем текст на слова
+    words = text.split()
+
+    # Лемматизируем каждое слово и убираем стоп-слова
+    lemmatized_words = [morph.parse(word)[0].normal_form for word in words]
+
+    # Собираем текст из лемматизированных слов
+    lemmatized_text = ' '.join(lemmatized_words)
+    return lemmatized_text
+
+
+def data_preprocessing(text):
+    # From Phase 1
+    text = re.sub(r':[a-zA-Z]+:', '', text)  # Убираем смайлики
+    text = text.lower()  # Переводим текст в нижний регистр
+    text = re.sub(r'@[\w_-]+', '', text)  # Убираем упоминания пользователей
+    text = re.sub(r'#(\w+)', '', text)  # Убираем хэштеги
+    text = re.sub(r'\d+', '', text)  # Убираем цифры
+    # Убираем ссылки
+    text = re.sub(
+        r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', '', text)
+    text = re.sub(r'\s+', ' ', text)  # Убираем лишние пробелы
+    # Удаление английских слов
+    text = ' '.join(re.findall(r'\b[а-яА-ЯёЁ]+\b', text))
+    # From Phase 2
+    text = re.sub("<.*?>", "", text)  # html tags
+    text = "".join([c for c in text if c not in string.punctuation])
+    splitted_text = [word for word in text.split() if word not in stop_words]
+    text = " ".join(splitted_text)
+    return text.strip()
+
+
+def preprocess_single_string(
+    input_string: str,
+    seq_len: int,
+    vocab_to_int: dict,
+    verbose: bool = False
+) -> Tensor:
+    """Function for all preprocessing steps on a single string
+
+    Args:
+        input_string (str): input single string for preprocessing
+        seq_len (int): max length of sequence, it len(review_int[i]) > seq_len it will be trimmed, else it will be padded by zeros
+        vocab_to_int (dict, optional): word corpus {'word' : int index}. Defaults to vocab_to_int.
+
+    Returns:
+        list: preprocessed string
+    """
+    preprocessed_string = lemmatize(input_string)
+    preprocessed_string = data_preprocessing(input_string)
+    result_list = []
+    for word in preprocessed_string.split():
+        try:
+            result_list.append(vocab_to_int[word])
+        except KeyError as e:
+            if verbose:
+                print(f'{e}: not in dictionary!')
+            pass
+    result_padded = padding([result_list], seq_len)[0]
+
+    return Tensor(result_padded)
+# ------------------------------------------------------------#
+
+
+st.title("Классификация тематики новостей из телеграм каналов")
+# st.write('Model summary:')
+text = st.text_input('Input some news')
+text_4_test = text
+
+# Загрузка словаря из файла
+with open('model/vocab_to_int.pkl', 'rb') as f:
+    vocab_to_int = pickle.load(f)
+
+if text != '':
+    test_review = preprocess_single_string(
+        text_4_test, net_config.seq_len, vocab_to_int)
+    test_review = torch.tensor(test_review, dtype=torch.int64)
+    result = torch.sigmoid(model_lstm(test_review.unsqueeze(0)))
+    num = result.argmax().item()
+
+st.write('---')
+st.write('Initial text:')
+st.write(text)
+st.write('---')
+st.write('Preprocessing:')
+st.write(data_preprocessing(text))
+st.write('---')
+st.write('Classes:')
+classes = ['крипта', 'мода', 'спорт', 'технологии', 'финансы']
+st.write('крипта *', 'мода *', 'спорт *', 'технологии *', 'финансы')
+st.write('---')
+
+st.write('Predict:')
+if text != '':
+    st.write('Classification: ', classes[num])
+    st.write('Label num: ', num)
+
+# Загружаем изображение через PIL
+image = Image.open("images/tg_metrics.png")
+
+# Отображение
+st.image(image, caption="Кошмареус переобучения", use_column_width=True)