from lightfm import LightFM
from lightfm.datasets import fetch_movielens
if __name__ == "__main__":
no_components = int(sys.argv[1])
print(f"Number of components is set to {no_components}")
# Load the MovieLens 100k dataset. Only five
# star ratings are treated as positive.
data = fetch_movielens(min_rating=5.0)
# Instantiate and train the model
model = LightFM(no_components=no_components, loss='warp')
model.fit(data['train'], epochs=30, num_threads=2)
joblib.dump(model, "model.joblib")