[ML] TfidfVectorizer를 사용한 Train, Test set 생성

Train Set 생성

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

# TfidfVectorizer 객체 생성
vectorizer = TfidfVectorizer()

# TfidfVectorizer 할 X_train 데이터
X_train = ['I go to my home my home is very large', 'I went out my home I go to the market']

# vectorizer는 X_train에 등장하는 단어를 통해 단어 사전을 생성한 뒤,
# 해당 단어 사전을 사용해 X_train를 TfidfVectorizer 한다.
# 따라서, vec_X_train의 차원은 (X_train 문장 수, 사전의 단어 수) 가 된다.
vec_X_train = vectorizer.fit_transform(X_train)

# 벡터라이저가 학습한 단어사전 정렬
word_list = [i[0] for i in sorted(vectorizer.vocabulary_.items())]

# vec_X_train을 DataFrame 화
pd.DataFrame(vec_X_train.toarray(), columns=[word_list])

Test Set 생성

# TfidfVectorizer 할 X_test 데이터
X_test = ['I bought a yellow lemon I go back to home']

# X_train TfidfVectorizer 시 사용했던 vectorizer 를 사용해, 
# X_test를 TfidfVectorizer 한다.
vec_X_test = vectorizer.transform(X_test)

# vec_X_test을 DataFrame 화
pd.DataFrame(vec_X_test.toarray(), columns=[word_list])