Clean up generation code