from gensim.models import Word2Vec
from nltk.cluster import KMeansClusterer
import nltk
import numpy as np 
 
sentences = [['We', 'should', 'watch', 'a', 'movie'],
            ['Babe', 'Ruth',  'was', 'a', 'great', 'baseball', 'player'],
            ['Lou', 'Gerhig', 'played', 'baseball'],
            ['Do', 'not', 'discuss', 'politics', 'at', 'work'],
            ['Baseball', 'hotdogs', 'Apple', 'Pie', 'and', 'Chevrolet'],
            ['Data', 'mining', 'can', 'use', 'machine', 'learning'],
            ['Clustering', 'uses', 'unsupervised', 'machine', 'learning'],
            ['My', 'company', 'does', 'machine', 'learning'],
            ['Bill', 'Gates', 'was', 'a', 'programmer'],  
            ['The', 'movie', 'was', 'bad']]
  
model = Word2Vec(sentences, min_count=1)
 
Data = []

for sentence in sentences:
    vector = []
    wordCount = 0
    for word in sentence:
       if wordCount == 0:
          vector = model[word]
       else:
          vector = np.add(vector, model[word])
       wordCount += 1
     
    Data.append(np.asarray(vector)/wordCount)
 
km = KMeansClusterer(5, nltk.cluster.util.euclidean_distance, repeats=1000)
assigned_clusters = km.cluster(Data, assign_clusters=True)
  
for index, sentence in enumerate(sentences):    
    print (str(assigned_clusters[index]) + ":" + str(sentence))
