#!/usr/bin/env python
# coding: utf-8

# # sklearn 如何计算 TFIDF
# 
# 此 notebook 是博文 [sklearn 如何计算 TFIDF](https://alanlee.fun/2019/08/09/sklearn-tfidf/) 中的代码。

# In[1]:


import numpy as np
import jieba
from sklearn.feature_extraction.text import TfidfTransformer, TfidfVectorizer, CountVectorizer


# In[2]:


documents = [
    "低头亲吻我的左手",
    "换取被宽恕的承诺",
    "老旧管风琴在角落",
    "一直一直一直伴奏",
]


# In[3]:


documents = [" ".join(jieba.cut(item)) for item in documents]


# In[4]:


documents


# ## 手算
# 
# 以 `一直` 为例

# In[5]:


# idf(一直, D)
np.log((1+4)/(1+1)) + 1


# In[6]:


# tfidf(一直, 文档 4, D)
3 * 1.916290731874155


# In[7]:


# 文档 4 未归一化 tfidf 向量
no_norm = np.array([5.748872195622465, 0, 1.916290731874155, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])


# In[8]:


# 文档 4 归一化 tfidf 向量
normed = no_norm / np.sqrt(sum(no_norm ** 2))
normed


# ## 使用 sklearn 计算

# In[9]:


# 默认情况下 sklearn 会莫名其妙地去除掉一些停用词，即使 stop_words=None 
# 详细讨论参见 https://github.com/scikit-learn/scikit-learn/issues/10756
vectorizer = TfidfVectorizer(token_pattern=r'(?u)\b\w+\b', strip_accents=False)
X = vectorizer.fit_transform(documents)


# In[10]:


vectorizer.vocabulary_


# In[11]:


vectorizer.idf_


# In[12]:


X.toarray()


# In[13]:


vectorizer.stop_words_