wehub-resource-sync
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""
|
|
Trains an Alternating Least Squares (ALS) model for user/movie ratings.
|
|
The input is a Parquet ratings dataset (see etl_data.py), and we output
|
|
an mlflow artifact called 'als-model'.
|
|
"""
|
|
|
|
import click
|
|
import pyspark
|
|
from pyspark.ml import Pipeline
|
|
from pyspark.ml.evaluation import RegressionEvaluator
|
|
from pyspark.ml.recommendation import ALS
|
|
|
|
import mlflow
|
|
import mlflow.spark
|
|
|
|
|
|
@click.command()
|
|
@click.option("--ratings-data")
|
|
@click.option("--split-prop", default=0.8, type=float)
|
|
@click.option("--max-iter", default=10, type=int)
|
|
@click.option("--reg-param", default=0.1, type=float)
|
|
@click.option("--rank", default=12, type=int)
|
|
@click.option("--cold-start-strategy", default="drop")
|
|
def train_als(ratings_data, split_prop, max_iter, reg_param, rank, cold_start_strategy):
|
|
seed = 42
|
|
|
|
with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
|
|
ratings_df = spark.read.parquet(ratings_data)
|
|
(training_df, test_df) = ratings_df.randomSplit([split_prop, 1 - split_prop], seed=seed)
|
|
training_df.cache()
|
|
test_df.cache()
|
|
|
|
mlflow.log_metric("training_nrows", training_df.count())
|
|
mlflow.log_metric("test_nrows", test_df.count())
|
|
|
|
print(f"Training: {training_df.count()}, test: {test_df.count()}")
|
|
|
|
als = (
|
|
ALS()
|
|
.setUserCol("userId")
|
|
.setItemCol("movieId")
|
|
.setRatingCol("rating")
|
|
.setPredictionCol("predictions")
|
|
.setMaxIter(max_iter)
|
|
.setSeed(seed)
|
|
.setRegParam(reg_param)
|
|
.setColdStartStrategy(cold_start_strategy)
|
|
.setRank(rank)
|
|
)
|
|
|
|
als_model = Pipeline(stages=[als]).fit(training_df)
|
|
|
|
reg_eval = RegressionEvaluator(
|
|
predictionCol="predictions", labelCol="rating", metricName="mse"
|
|
)
|
|
|
|
predicted_test_dF = als_model.transform(test_df)
|
|
|
|
test_mse = reg_eval.evaluate(predicted_test_dF)
|
|
train_mse = reg_eval.evaluate(als_model.transform(training_df))
|
|
|
|
print(f"The model had a MSE on the test set of {test_mse}")
|
|
print(f"The model had a MSE on the (train) set of {train_mse}")
|
|
mlflow.log_metric("test_mse", test_mse)
|
|
mlflow.log_metric("train_mse", train_mse)
|
|
mlflow.spark.log_model(als_model, artifact_path="als-model")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
train_als()
|