updated get_prediction to be not fucked