File size: 2,860 Bytes
431bcf6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import os
import pandas as pd
import numpy as np
import joblib
from sklearn.linear_model import RidgeCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
def evaluate_model(model, X, y, name):
preds = model.predict(X)
rmse = np.sqrt(mean_squared_error(y, preds))
mae = mean_absolute_error(y, preds)
r2 = r2_score(y, preds)
print(f"{name} Metrics:")
print(f" RMSE: {rmse:.4f}")
print(f" MAE: {mae:.4f}")
print(f" R2: {r2:.4f}\n")
return rmse, mae, r2
def main():
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
data_dir = os.path.join(base_dir, "data")
models_dir = os.path.join(base_dir, "models")
os.makedirs(models_dir, exist_ok=True)
# ---------------------------------------------------------
# 1. Classical Ridge Model
# ---------------------------------------------------------
print("--- Training Classical Ridge Model ---")
train_feat = pd.read_csv(os.path.join(data_dir, "train_features.csv"))
val_feat = pd.read_csv(os.path.join(data_dir, "validation_features.csv"))
test_feat = pd.read_csv(os.path.join(data_dir, "test_features.csv"))
y_train = train_feat['extraversion'].values
X_train = train_feat.drop(columns=['extraversion']).values
y_val = val_feat['extraversion'].values
X_val = val_feat.drop(columns=['extraversion']).values
y_test = test_feat['extraversion'].values
X_test = test_feat.drop(columns=['extraversion']).values
# Train
classical_ridge = RidgeCV(alphas=np.logspace(-3, 3, 7), cv=5)
classical_ridge.fit(X_train, y_train)
# Evaluate
evaluate_model(classical_ridge, X_val, y_val, "Classical Ridge (Validation)")
evaluate_model(classical_ridge, X_test, y_test, "Classical Ridge (Test)")
joblib.dump(classical_ridge, os.path.join(models_dir, "classical_ridge_model.pkl"))
print("Saved classical_ridge_model.pkl\n")
# ---------------------------------------------------------
# 2. BERT Ridge Model
# ---------------------------------------------------------
print("--- Training BERT Ridge Model ---")
train_bert = np.load(os.path.join(data_dir, "train_bert_embeddings.npy"))
val_bert = np.load(os.path.join(data_dir, "validation_bert_embeddings.npy"))
test_bert = np.load(os.path.join(data_dir, "test_bert_embeddings.npy"))
# Train
bert_ridge = RidgeCV(alphas=np.logspace(-3, 3, 7), cv=5)
bert_ridge.fit(train_bert, y_train)
# Evaluate
evaluate_model(bert_ridge, val_bert, y_val, "BERT Ridge (Validation)")
evaluate_model(bert_ridge, test_bert, y_test, "BERT Ridge (Test)")
joblib.dump(bert_ridge, os.path.join(models_dir, "bert_ridge_model.pkl"))
print("Saved bert_ridge_model.pkl\n")
if __name__ == "__main__":
main()
|