IS_445_Observatory / train_model.py
KeegBarb's picture
Update train_model.py
f24417e verified
import pandas as pd
import numpy as np
import joblib
from sklearn.linear_model import Ridge
from sklearn.model_selection import train_test_split
data_url = "https://raw.githubusercontent.com/KeeganBarbee/KeeganBarbee.github.io/main/OnlineNewsPopularity.csv"
df = pd.read_csv(data_url)
df.columns = df.columns.str.strip()
df['log_shares'] = np.log1p(df['shares'])
feature_cols = ['n_tokens_content', 'num_imgs', 'global_sentiment_polarity', 'global_subjectivity', 'title_sentiment_polarity',
'weekday_is_monday', 'weekday_is_tuesday', 'weekday_is_wednesday', 'weekday_is_thursday', 'weekday_is_friday', 'weekday_is_saturday',
'weekday_is_sunday', 'n_tokens_title', 'num_videos', 'num_hrefs']
X = df[feature_cols]
y = df['log_shares']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model = Ridge(alpha=1.0)
model.fit(X_train, y_train)
joblib.dump(model, 'popularity_model.pkl')
joblib.dump(feature_cols, 'model_features.pkl')