From df7658672fb4a140dd0f3c561901a9a158436b7c Mon Sep 17 00:00:00 2001
From: Tony Reina <37851530+tonyreina@users.noreply.github.com>
Date: Fri, 28 May 2021 11:43:58 -0700
Subject: [PATCH 1/2] Update features.py

---
 mol2vec/features.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/mol2vec/features.py b/mol2vec/features.py
index 3ce0409..e869e7e 100644
--- a/mol2vec/features.py
+++ b/mol2vec/features.py
@@ -15,6 +15,9 @@
 import timeit
 from joblib import Parallel, delayed
 
+import gensim
+
+GENSIM4 = gensim.__version__ > "4"  # GENSIM4 has some API changes from GENSIM 3
 
 class DfVec(object):
     """
@@ -422,7 +425,10 @@ def sentences2vec(sentences, model, unseen=None):
     -------
     np.array
     """
-    keys = set(model.wv.vocab.keys())
+   if GENSIM4:
+        keys = set(model.wv.key_to_index)
+    else:
+        keys = set(model.wv.vocab.keys())
     vec = []
     if unseen:
         unseen_vec = model.wv.word_vec(unseen)

From abd771e5a1a444f25283d218aa44862201981b4f Mon Sep 17 00:00:00 2001
From: Tony Reina <37851530+tonyreina@users.noreply.github.com>
Date: Sat, 29 May 2021 15:41:00 -0700
Subject: [PATCH 2/2] Update features.py

---
 mol2vec/features.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mol2vec/features.py b/mol2vec/features.py
index e869e7e..c192ab8 100644
--- a/mol2vec/features.py
+++ b/mol2vec/features.py
@@ -425,7 +425,7 @@ def sentences2vec(sentences, model, unseen=None):
     -------
     np.array
     """
-   if GENSIM4:
+    if GENSIM4:
         keys = set(model.wv.key_to_index)
     else:
         keys = set(model.wv.vocab.keys())