计算字符串的tf-idf

Question

我有两个文件，分别是 doc1.txt 和 doc2.txt。这两个文件的内容是：

 #doc1.txt
 very good, very bad, you are great

 #doc2.txt
 very bad, good restaurent, nice place to visit

我想把我的文本内容用 , 逗号分开，这样我的最终结果 DocumentTermMatrix 就会变成：

      terms
 docs       very good      very bad        you are great   good restaurent   nice place to visit
 doc1       tf-idf          tf-idf         tf-idf          0                    0
 doc2       0               tf-idf         0               tf-idf             tf-idf

我知道怎么计算单个单词的 DocumentTermMatrix（可以参考 这个链接），但是我不知道怎么在 Python 中计算 字符串 的 DocumentTermMatrix。

Answer 1

你可以给TfidfVectorizer这个工具指定一个analyzer参数，来用你自己的方式提取特征：

from sklearn.feature_extraction.text import TfidfVectorizer

docs = ['very good, very bad, you are great',
        'very bad, good restaurent, nice place to visit']

tfidf = TfidfVectorizer(analyzer=lambda d: d.split(', ')).fit(docs)
print tfidf.get_feature_names()

这样得到的特征是：

['good restaurent', 'nice place to visit', 'very bad', 'very good', 'you are great']

---

如果你真的无法一次性把所有数据都加载到内存中，可以使用这个方法：

from sklearn.feature_extraction.text import TfidfVectorizer

docs = ['doc1.txt', 'doc2.txt']

def extract(filename):
    with open(filename) as f:
        features = []
        for line in f:
            features += line.strip().split(', ')
        return features

tfidf = TfidfVectorizer(analyzer=extract).fit(docs)
print tfidf.get_feature_names()

这个方法会一次加载一个文档，而不是一次性把所有文档都放到内存里。