prediction function updated