How to populate the xticklabels with vocabulary in seaborn heatmap?

user_649994 0
Viewed 220

Written a small code snippet to display collapsed one-hot representation of my sample corpus data

corpus = ['Time flies like an arrow.', 'Fruit flies like a banana.']

I want to print vocabulary in the x-axis heatmap below instead of numbering. i.e. an, arrow, banana, flies, fruit, like, time intead of 1, 2,3,4,5,6 in x-axis

Here is the output of current heatmap.

enter image description here

Currently I use xticklabels='auto'. What should I use here to print the unique vocabulary instead?

Here is my code snippet

from sklearn.feature_extraction.text import CountVectorizer

import seaborn as sns

corpus = ['Time flies like an arrow.',
          'Fruit flies like a banana.']

one_hot_vectorizer = CountVectorizer(binary=True)

one_hot = one_hot_vectorizer.fit_transform(corpus).toarray()

sns.heatmap(one_hot, annot=True, cbar = False, xticklabels='auto', yticklabels=['Sentence1','Sentence 2'])

Thanks!

scikit-learn seaborn heatmap python-3.x python
1 Answers

You can extract the names from your vectorizer using one_hot_vectorizer.get_feature_names() and then pass it to the xticklabels kwargs:

from sklearn.feature_extraction.text import CountVectorizer

import seaborn as sns

corpus = ['Time flies like an arrow.',
          'Fruit flies like a banana.']

one_hot_vectorizer = CountVectorizer(binary=True)

one_hot = one_hot_vectorizer.fit_transform(corpus).toarray()

xticklabels = one_hot_vectorizer.get_feature_names()
sns.heatmap(one_hot, annot=True, cbar = False, xticklabels=xticklabels, yticklabels=['Sentence1','Sentence 2'])

Output:

enter image description here

user_105070360
Related