Python - Using TF-IDF to summarise dataframe text column

Viewed 409

I have a dataframe with a column containing text.

I want to create a new column that contains a tuple/list of the top 'n' TF-IDF scoring words in each row as a way of summarizing what is in the text.

An example dataframe (with a large amount of brevity) is:

df = pd.DataFrame({'Ref': [1,2,3,4,5], 'Text': ["the cow jumped off the other cow", 
                                                "the fox had a fox", 
                                                "the spanner was a tool to tool", 
                                                "the football player played football",
                                                "the house had a house"]})

I have spent the last few days trying to find a solution, but I can only find examples which finds the top TF-IDF words for the whole corpus, rather than for each row in a dataframe based on the whole corpus.

Can anyone steer me in the right direction?

1 Answers

Here is a possible solution:

from sklearn.feature_extraction.text import TfidfVectorizer
import numpy as np
import pandas as pd

n = 3 # top n TF-IDF words

tfidf = TfidfVectorizer(token_pattern=r"\w+") # no words are left out
X = tfidf.fit_transform(df['Text'])
ind = (-X.todense()).argpartition(n)[:, :n]
top_words = pd.Series(
    map(
        lambda words_values: dict(zip(*words_values)),
        zip(
            np.array(tfidf.get_feature_names())[ind],
            np.asarray(np.take_along_axis(X, ind, axis=1).todense()),
        ),
    ),
)

Here is the result:

>>> top_words
0    {'cow': 0.7111977362687212, 'other': 0.3555988681343606, 'off': 0.3555988681343606}
1    {'fox': 0.8665817814049075, 'had': 0.34957636239744133, 'a': 0.2901799593148741}
2    {'tool': 0.7218960199361867, 'was': 0.36094800996809334, 'spanner': 0.36094800996809334}
3    {'football': 0.8014723840888909, 'player': 0.40073619204444544, 'played': 0.40073619204444544}
4    {'house': 0.8665817814049075, 'had': 0.34957636239744133, 'a': 0.2901799593148741}
Related