import nltk
from gensim.models import KeyedVectors
import pickle

embeddings = KeyedVectors.load_word2vec_format('GoogleNews-vectors-negative300.bin.gz', binary=True)
print('opened Google vectors file')
f = open('capitals.txt', 'r').read()
set_words = set(nltk.word_tokenize(f))
select_words = words = ['king', 'queen', 'oil', 'gas', 'happy', 'sad', 'city', 'town', 'village', 'country',
                        'continent', 'petroleum', 'joyful']
print('adds words in set')

for w in select_words:
    set_words.add(w)


def get_word_embeddings(embeddings):
    word_embeddings = {}
    for word in embeddings.vocab:
        if word in set_words:
            word_embeddings[word] = embeddings[word]
    return word_embeddings


# Testing your function
word_embeddings = get_word_embeddings(embeddings)
print(len(word_embeddings))
pickle.dump(word_embeddings, open("word_embeddings_subset.p", "wb"))
