Update tfidf.py

tfidf.py

@@ -1,4 +1,7 @@
-
+from sklearn.model_selection import train_test_split
+X = df['title']
+y = df['labels']
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 from sklearn.feature_extraction.text import TfidfVectorizer
 tfidf = TfidfVectorizer(max_features=5000, ngram_range=(1, 2), stop_words='english')
 X_train_tfidf = tfidf.fit_transform(X_train)