pythonpandasgensim

How to turn embeddings loaded in a Pandas DataFrame into a Gensim model?


I have a DataFrame in which the index are words and I have 100 columns with float number such that for each word I have its embedding as a 100d vector. I would like to convert my DataFrame object into a gensim model object so that I can use its methods; specially gensim.models.keyedvectors.most_similar() so that I can search for similar words within my subset.

Which is the preferred way of doing that?

Thanks


Solution

  • Not sure what the "preferred" way of doing this is, but the format gensim expects is pretty easy to replicate:

    data = pd.DataFrame([[0.15941701, 0.84058299],
                         [0.12190033, 0.87809967],
                         [0.06293788, 0.93706212]],
                        index=["these", "be", "words"])
    
    np.savetxt('test.txt', data.reset_index().values, 
               delimiter=" ", 
               header="{} {}".format(len(data), len(data.columns)),
               comments="",
               fmt=["%s"] + ["%.18e"]*len(data.columns))
    

    The header is 2 space separated integers, the number of words in the vocabulary and the length of the word vector. The first column of each row is the word itself. The rest of the columns are the elements of the word vector. The fmt weirdness is to have the first element formatted as a string, and the rest formatted as a float.

    Then can load this in gensim and do whatever:

    import gensim
    
    from gensim.models.keyedvectors import KeyedVectors
    word_vectors = KeyedVectors.load_word2vec_format('test.txt', binary=False)
    
    word_vectors.similarity('these', 'words')