import nltk
from nltk.corpus import brown
from gensim.models import Word2Vec
import numpy as np
import faiss

nltk.download('brown')

corpus = [[word.lower() for word in sentence] for sentence in brown.sents()]

model = Word2Vec(
    corpus,
    vector_size=50,
    window=5,
    min_count=5,
    seed=42,
    workers=1
)

words = list(model.wv.index_to_key)
vectors = np.array(model.wv.vectors).astype("float32")

index = faiss.IndexFlatL2(50)
index.add(vectors)

query = "man"
query_vector = np.array([model.wv[query]]).astype("float32")

distances, indices = index.search(query_vector, 4)

print("Words similar to:", query)

for i in indices[0][1:]:
    print(words[i])