| import pandas as pd | |
| import joblib | |
| df = pd.read_csv("hf://datasets/CIS5190abcd/headlines_train/train_cleaned_headlines.csv") | |
| 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) | |
| X_test_tfidf = tfidf.transform(X_test) | |
| joblib.dump(X_train_tfidf, 'X_train_tfidf.pkl') | |
| joblib.dump(X_test_tfidf, 'X_test_tfidf.pkl') | |
| joblib.dump(y_train, 'y_train.pkl') | |
| joblib.dump(y_test, 'y_test.pkl') | |