Added step_callback, fixed bug w/multiple tokens