Update pages/imdb.py

pages/imdb.py

@@ -1,110 +1,157 @@
+import os
+import numpy as np
 import pandas as pd
+
+import matplotlib.pyplot as plt
 import streamlit as st
-import torch
+import re
+import string
+from collections import Counter
+
+from transformers import AutoConfig, AutoTokenizer, AutoModel, AutoModelForSequenceClassification, Trainer, TrainingArguments
+
+from gensim.models import Word2Vec
+from string import punctuation
 import transformers
+import warnings
+warnings.filterwarnings('ignore')
+
+from sklearn.model_selection import train_test_split
 import time
-import pickle
-import numpy as np
-from gensim.models import Word2Vec
-from sklearn.feature_extraction.text import TfidfVectorizer
+
+from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
+import pickle
+import torch
+from torch.utils.data import DataLoader, TensorDataset
 import torch.nn as nn
+import torchutils as tu
+from torchmetrics.classification import BinaryAccuracy
 from data.rnn_preprocessing import (
-    data_preprocessing,
-    preprocess_single_string
-)
-
-# Load Word2Vec model
-wv = Word2Vec.load('models/word2vec32.model')
-embedding_matrix = np.zeros((VOCAB_SIZE, EMBEDDING_DIM))
-vocab_to_int = {word: idx + 1 for idx, word in enumerate(wv.wv.index_to_key)}
-for word, i in vocab_to_int.items():
-    try:
-        embedding_vector = wv.wv[word]
-        embedding_matrix[i] = embedding_vector
-    except KeyError:
-        pass
-
-
-
-# Load LSTM model
-embedding_layer32 = nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))
-VOCAB_SIZE = len(vocab_to_int) + 1  # add 1 for the padding token
-HIDDEN_DIM = 64
-SEQ_LEN = 32
-
-
-
-
-# Load TF-IDF model
-tfidf_model = pickle.load(open('models/modeltfidf.sav', 'rb'))
-
-
-
-class LSTMClassifierBi32(nn.Module):
-    def __init__(self, embedding_dim: int, hidden_size: int = 32) -> None:
-        super().__init__()
-
-        self.embedding_dim = embedding_dim
-        self.hidden_size = hidden_size
-        self.embedding = embedding_layer32
-        self.lstm = nn.LSTM(
-            input_size=self.embedding_dim,
-            hidden_size=self.hidden_size,
-            batch_first=True,
-            bidirectional=True
-        )
-        self.clf = nn.Sequential(
-            nn.Linear(self.hidden_size * 2, 128),
-            nn.Dropout(),
-            nn.Sigmoid(),
-            nn.Linear(128, 64),
-            nn.Dropout(),
-            nn.Sigmoid(),
-            nn.Linear(64, 1)
-        )
-
-    def forward(self, x):
-        embeddings = self.embedding(x)
-        out, (_, _) = self.lstm(embeddings)
-        out = self.clf(out[:, -1, :])
-        return out
-
-
-model = LSTMClassifierBi32(embedding_dim=EMBEDDING_DIM, hidden_size=HIDDEN_DIM)
-model.load_state_dict(torch.load('models/ltsm_bi1.pt'))
-model.eval()
-
-def predict_sentence(text: str, model: nn.Module):
-    result = model(preprocess_single_string(text, seq_len=SEQ_LEN, vocab_to_int=vocab_to_int).unsqueeze(0)).sigmoid().round().item()
-    return 'negative' if result == 0.0 else 'positive'
-
+                                data_preprocessing, 
+                                preprocess_single_string
+                                )
 
 def main():
+    device = 'cpu'
     df = pd.read_csv('data/imdb.csv')
     df['sentiment'] = df['sentiment'].apply(lambda x: 1 if x == 'positive' else 0)
     reviews = df['review'].tolist()
     preprocessed = [data_preprocessing(review) for review in reviews]
 
-    tfid_vectorizer = TfidfVectorizer(max_df=0.5, min_df=5)
+    wv = Word2Vec.load('models/word2vec32.model')
+
+    words_list = [word for review in preprocessed for word in review.lower().split()]
+    for i in words_list:
+        ''.join([j for j in i if j not in punctuation])
+        
+    # делаем множество уникальных слов.
+    unique_words = set(words_list)
+
+    # word -> index
+    vocab_to_int = {word: idx+1 for idx, word in enumerate(sorted(unique_words))}
+
+    word_seq = [i.split() for i in preprocessed]
+    VOCAB_SIZE = len(vocab_to_int) + 1  # add 1 for the padding token
+    EMBEDDING_DIM = 32
+    HIDDEN_DIM = 64
+    SEQ_LEN = 32
+
+    embedding_matrix = np.zeros((VOCAB_SIZE, EMBEDDING_DIM))
+
+    for word, i in vocab_to_int.items():
+        try:
+            embedding_vector = wv.wv[word]
+            embedding_matrix[i] = embedding_vector
+        except KeyError:
+            pass
+
+    embedding_layer32 = torch.nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))
+
+
+    class LSTMClassifierBi32(nn.Module):
+        def __init__(self, embedding_dim: int, hidden_size:int = 32) -> None:
+            super().__init__()
+
+            self.embedding_dim = embedding_dim
+            self.hidden_size = hidden_size
+            self.embedding = embedding_layer32
+            self.lstm = nn.LSTM(
+                input_size=self.embedding_dim,
+                hidden_size=self.hidden_size,
+                batch_first=True,
+                bidirectional=True
+            )
+            self.clf = nn.Sequential(nn.Linear(self.hidden_size*2, 128), 
+                nn.Dropout(),
+                nn.Sigmoid(),
+                nn.Linear(128, 64),
+                nn.Dropout(),
+                nn.Sigmoid(),
+                nn.Linear(64, 1)
+            )
+
+        def forward(self, x):
+            embeddings = self.embedding(x)
+            out, (_, _) = self.lstm(embeddings)
+            out = self.clf(out[:,-1,:])
+            return out
+        
+    model = LSTMClassifierBi32(embedding_dim=EMBEDDING_DIM, hidden_size=HIDDEN_DIM)
+    model.load_state_dict(torch.load('models/ltsm_bi1.pt'))
+    model.eval()
+
+    def predict_sentence(text:str, model: nn.Module):
+        result = model.to(device)(preprocess_single_string(text, seq_len=SEQ_LEN, vocab_to_int=vocab_to_int).unsqueeze(0)).sigmoid().round().item()
+        return 'negative' if result == 0.0 else 'positive'
+    
+    #Bag Tfidf
+    # bagvectorizer = CountVectorizer(max_df=0.5,
+    # min_df=5,
+    # stop_words="english",)
+    # bvect = bagvectorizer.fit(preprocessed)
+    # X_bag = bvect.transform(preprocessed)
+
+    tfid_vectorizer = TfidfVectorizer(
+    max_df=0.5,
+    min_df=5)
     vect = tfid_vectorizer.fit(preprocessed)
     X_tfidf = vect.transform(preprocessed)
-
+    
+    tfidf_model = pickle.load(open('models/modeltfidf.sav', 'rb'))
+    # bag_model = pickle.load(open('models/modelbag.sav', 'rb'))
+    # def predictbag(text):
+    #     result = bag_model.predict(vect.transform([text]))
+    #     return 'negative' if result == [0] else 'positive'
+
+    def predicttf(text):
+        result = tfidf_model.predict(vect.transform([text]))
+        return 'negative' if result == [0] else 'positive'
+    
+        
+
+
+    
+    
+    
+    
+    
     review = st.text_input('Enter review')
 
     start1 = time.time()
-
+    
+    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+    config = AutoConfig.from_pretrained('distilbert-base-uncased', num_labels=2)
+    
+    automodel = AutoModelForSequenceClassification.from_config(config)
     autotoken = transformers.AutoTokenizer.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english')
 
     input_tokens = autotoken(
-        review,
-        return_tensors='pt',
-        padding=True,
+        review, 
+        return_tensors='pt', 
+        padding=True, 
         max_length=10
     )
-
-    config = transformers.AutoConfig.from_pretrained('distilbert-base-uncased', num_labels=2)
-    automodel = transformers.AutoModelForSequenceClassification.from_config(config)
     outputs = automodel(**input_tokens)
     st.write('Sentiment Predictions')
     st.write(f'\nBERT: {[automodel.config.id2label[i.item()] for i in outputs.logits.argmax(-1)]}')
@@ -112,20 +159,20 @@ def main():
     st.write(f'{(end1 - start1):.2f} sec')
     start2 = time.time()
 
-    st.write(f'LSTM: {predict_sentence(review, model)}')
+    st.write(f'LTSM: {predict_sentence(review, model)}')
     end2 = time.time()
     st.write(f'{(end2 - start2):.2f} sec')
-
+    # start3 = time.time()
+    # st.write(f'bag+log: {predictbag(review)}')
+    # end3 = time.time()
+    # st.write(f'{(end3 - start3):.2f} sec')
     start4 = time.time()
-    st.write(f'TF-IDF+Logistic Regression: {predicttf(review)}')
+    st.write(f'tfidf+log: {predicttf(review)}')
     end4 = time.time()
     st.write(f'{(end4 - start4):.2f} sec')
 
 
-def predicttf(text):
-    result = tfidf_model.predict(vect.transform([text]))
-    return 'negative' if result == [0] else 'positive'
-
+    
 
 if __name__ == '__main__':
-    main()
+    main()