Fix UX in generator