Updating generation code. Closes #9