在scikit-learn的tf-idf矩阵中获取文档名称

Question

我已经创建了一个tf-idf矩阵，但现在我想检索每篇文档的前2个词语。 我想传递文档ID并获得前2个词语。

目前，我有这个样本数据：

from sklearn.feature_extraction.text import TfidfVectorizer

d = {'doc1':"this is the first document",'doc2':"it is a sunny day"} ### corpus

test_v = TfidfVectorizer(min_df=1)    ### applied the model
t = test_v.fit_transform(d.values())
feature_names = test_v.get_feature_names() ### list of words/terms

>>> feature_names
['day', 'document', 'first', 'is', 'it', 'sunny', 'the', 'this']

>>> t.toarray()
array([[ 0.        ,  0.47107781,  0.47107781,  0.33517574,  0.        ,
     0.        ,  0.47107781,  0.47107781],
   [ 0.53404633,  0.        ,  0.        ,  0.37997836,  0.53404633,
     0.53404633,  0.        ,  0.        ]])

我可以通过提供行号来访问矩阵，例如：

 >>> t[0,1]
   0.47107781233161794

有没有一种方法可以通过文档ID访问此矩阵？在我的情况下是'doc1'和'doc2'。

谢谢

- user1525721

---

不是直接的，但你可以将数据封装在pandas DataFrame中。 - BrenBarn

在scikit-learn中没有“文档名称”的概念。您必须自己存储这些。 - Fred Foo

我也是这么想的。你们证实了它。感谢你们的建议。 - user1525721

Answer 1

通过执行

t = test_v.fit_transform(d.values())

你将失去与文档ID的任何链接。字典是无序的，因此你不知道哪个值对应哪个顺序。这意味着在将值传递给fit_transform函数之前，需要记录哪个值对应哪个ID。

例如，你可以这样做：

counter = 0
values = []
key = {}


for k,v in d.items():
    values.append(v)
    key[k] = counter
    counter+=1

t = test_v.fit_transform(values)

你可以构建一个函数来通过文档ID访问此矩阵：

def get_doc_row(docid):
    rowid = key[docid]
    row = t[rowid,:]
    return row

- patapouf_ai