adapted code to new model