predict the output in python