Update pages/imdb.py

pages/imdb.py

@@ -16,15 +16,21 @@ from data.rnn_preprocessing import (
 
 # Load Word2Vec model
 wv = Word2Vec.load('models/word2vec32.model')
-embedding_matrix = wv.wv.vectors
+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 TF-IDF model
 tfidf_model = pickle.load(open('models/modeltfidf.sav', 'rb'))
 
 # Load LSTM model
-embedding_layer32 = torch.nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))
-VOCAB_SIZE, EMBEDDING_DIM = embedding_matrix.shape
+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
 
@@ -63,7 +69,6 @@ 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'