Fix bug in getNearestNeighbors() from module_embedding.py

modules/model_embbeding.py

@@ -5,7 +5,6 @@ from sklearn.decomposition import PCA
 from gensim.models import KeyedVectors
 from typing import List
 import os
-import operator
 import pandas as pd
 
 import numpy as np
@@ -136,14 +135,12 @@ class Embedding:
         self, 
         word: str
     ):
-
         return self.__getValue(word, 'embedding')
 
     def getPCA(
         self, 
         word: str
     ):
-
         return self.__getValue(word, 'pca')
     
     def getNearestNeighbors(
@@ -160,8 +157,8 @@ class Embedding:
             
         elif nn_method == 'sklearn':
             word_emb = self.getEmbedding(word).reshape(1,-1)
-            _, nn_ids = self.neigh.kneighbors(word_emb, n_neighbors)
-            words = operator.itemgetter(*nn_ids[0])(self.ds['word'].to_list())
+            _, nn_ids = self.neigh.kneighbors(word_emb, n_neighbors + 1)      #Fix and Update
+            words = [self.ds['word'].to_list()[idx] for idx in nn_ids[0]][1:] #Fix and Update
         else: 
             words = []
         return words