add predict function