File size: 695 Bytes
c2e00ba 26f1a8a 16a82d8 f7ba0f9 fad1a71 9222ce0 fad1a71 26f1a8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
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=7000, ngram_range=(1, 3), 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')
|