위키피디아 영화 줄거리 데이터로 주제분석을 한다. 데이터를 다운받는다.
import requests
res = requests.get('https://github.com/euphoris/datasets/raw/master/movie-plot.zip')
with open('movie-plot.zip', 'wb') as f:
f.write(res.content)
판다스로 연다.
import pandas as pd
df = pd.read_csv('movie-plot.zip', index_col=0)
불러온 데이터를 확인한다.
df.head()
| Title | Plot | |
|---|---|---|
| 0 | Kansas Saloon Smashers | A bartender is working at a saloon, serving dr... |
| 1 | Love by the Light of the Moon | The moon, painted with a smiling face hangs ov... |
| 2 | The Martyred Presidents | The film, just over a minute long, is composed... |
| 3 | Terrible Teddy, the Grizzly King | Lasting just 61 seconds and consisting of two ... |
| 4 | Jack and the Beanstalk | The earliest known adaptation of the classic f... |
from sklearn.model_selection import train_test_split
index_train, index_test = train_test_split(df.index, test_size=.2, random_state=1234)
훈련 데이터셋
plot_train = df.loc[index_train, 'Plot']
테스트 데이터셋의 크기
plot_test = df.loc[index_test, 'Plot']
from sklearn.feature_extraction.text import CountVectorizer
cv = CountVectorizer(stop_words='english', max_features=2000)
훈련 데이터로 단어 문서 행렬을 만든다.
x_train = cv.fit_transform(plot_train)
테스트 데이터도 단어 문서 행렬도 만든다. 훈련 데이터와 같은 형식을 유지해야 하므로 transform을 사용한다.
x_test = cv.transform(plot_test)
이후의 분석을 위해 데이터를 저장한다.
import joblib
data = {
'x_train': x_train,
'x_test': x_test,
'index_train': index_train,
'index_test': index_test,
'cv': cv
}
joblib.dump(data, 'movie.pkl')
['movie.pkl']