import numpy as np
from gensim.models import Word2Vec
from sklearn.manifold import TSNE

modell = Word2Vec.load("w2vlemmas.model")
keys = ['Шекспир', 'Пушкин', 'Гоголь', 'матрос', 'кот', 'роман']
embedding_clusters = []
word_clusters = []
for word in keys:
    embeddings = []
    words = []
    for similar_word, _ in modell.wv.most_similar(word, topn=30):
        words.append(similar_word)
        embeddings.append(modell.wv[similar_word])
    embedding_clusters.append(embeddings)
    word_clusters.append(words)
tsne_model_en_2d = TSNE(perplexity=15, n_components=2, init='pca', n_iter=3500, random_state=32)
embedding_clusters = np.array(embedding_clusters)
n, m, k = embedding_clusters.shape
embeddings_en_2d = np.array(tsne_model_en_2d.fit_transform(embedding_clusters.reshape(n * m, k))).reshape(n, m, 2)