Small changes to finalize script functionality